feat: improvements on training and create new models and tools file

Author: renan-siqueira

.gitignore

@@ -3,6 +3,7 @@
 __pycache__
 
 samples
+params.json
 
 *.pth
 *.png

copy_randomic_files.py

@@ -0,0 +1,50 @@
+import os
+import shutil
+import random
+
+from settings.settings import (
+    COPY_DESTINATION_FOLDER,
+    COPY_SOURCE_FOLDER,
+    COPY_PERCENTAGE_TO_COPY,
+    COPY_RANDOM_MODE,
+    COPY_FIXED_NUMBER_TO_COPY
+)
+
+
+def copy_files(
+        source_folder,
+        destination_folder,
+        percentage=None,
+        fixed_number=None,
+        random_mode=True
+    ):
+
+    files = [file_name for file_name in os.listdir(source_folder)
+             if os.path.isfile(os.path.join(source_folder, file_name))]
+
+    if percentage is not None:
+        total_to_copy = int(len(files) * percentage / 100)
+    elif fixed_number is not None:
+        total_to_copy = min(fixed_number, len(files))
+    else:
+        raise ValueError("Either percentage or fixed_number must be provided!")
+
+    if random_mode:
+        chosen_files = random.sample(files, total_to_copy)
+    else:
+        chosen_files = files[:total_to_copy]
+
+    for file_name in chosen_files:
+        shutil.copy2(os.path.join(source_folder, file_name), destination_folder)
+
+    print(f"{total_to_copy} files have been copied from {source_folder} to {destination_folder}.")
+
+
+if __name__ == '__main__':
+    copy_files(
+        COPY_SOURCE_FOLDER,
+        COPY_DESTINATION_FOLDER,
+        percentage=COPY_PERCENTAGE_TO_COPY,
+        fixed_number=COPY_FIXED_NUMBER_TO_COPY,
+        random_mode=COPY_RANDOM_MODE
+    )

json/params.example.json

@@ -0,0 +1,9 @@
+{
+    "batch_size": 32,
+    "resolution": 64,
+    "encoding_dim": 16,
+    "num_epochs": 100,
+    "learning_rate": 0.001,
+    "ae_type": "ae",
+    "save_checkpoint": null
+}

json/params.json

@@ -2,8 +2,8 @@
     "batch_size": 32,
     "resolution": 64,
     "encoding_dim": 16,
-    "num_epochs": 25,
-    "learning_rate": 0.001,
-    "ae_type": "conv_dae",
+    "num_epochs": 10,
+    "learning_rate": 0.0001,
+    "ae_type": "vae",
     "save_checkpoint": null
 }

models/__init__.py

@@ -1,6 +1,8 @@
 from .autoencoder import Autoencoder
 from .autoencoder_dae import DenoisingAutoencoder
+from .autoencoder_sparse import SparseAutoencoder
 from .autoencoder_vae import VariationalAutoencoder
 from .convolutional_autoencoder import ConvolutionalAutoencoder
 from .convolutional_dae import DenoisingConvolutionalAutoencoder
 from .convolutional_vae import ConvolutionalVAE
+from .convolutional_sparse import SparseConvolutionalAutoencoder

models/autoencoder_sparse.py

@@ -0,0 +1,43 @@
+import torch.nn as nn
+
+
+class SparseAutoencoder(nn.Module):
+    def __init__(self, input_dim, encoding_dim):
+        super(SparseAutoencoder, self).__init__()
+
+        self.model_structure = 'linear'
+        self.model_variant = 'sparse'
+
+        self.encoder = nn.Sequential(
+            nn.Linear(input_dim, 1024),
+            nn.LeakyReLU(0.01),
+            nn.Linear(1024, 512),
+            nn.LeakyReLU(0.01),
+            nn.Linear(512, 256),
+            nn.LeakyReLU(0.01),
+            nn.Linear(256, 128),
+            nn.LeakyReLU(0.01),
+            nn.Linear(128, 64),
+            nn.LeakyReLU(0.01),
+            nn.Linear(64, encoding_dim)
+        )
+
+        self.decoder = nn.Sequential(
+            nn.Linear(encoding_dim, 64),
+            nn.ReLU(),
+            nn.Linear(64, 128),
+            nn.ReLU(),
+            nn.Linear(128, 256),
+            nn.ReLU(),
+            nn.Linear(256, 512),
+            nn.ReLU(),
+            nn.Linear(512, 1024),
+            nn.ReLU(),
+            nn.Linear(1024, input_dim),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        encoded = self.encoder(x)
+        decoded = self.decoder(encoded)
+        return decoded, encoded

models/convolutional_sparse.py

@@ -0,0 +1,41 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SparseConvolutionalAutoencoder(nn.Module):
+    def __init__(self):
+        super(SparseConvolutionalAutoencoder, self).__init__()
+
+        self.model_structure = 'convolutional'
+        self.model_variant = 'sparse'
+
+        # Encoder
+        self.enc0 = nn.Conv2d(3, 256, kernel_size=3, padding=1)
+        self.enc1 = nn.Conv2d(256, 128, kernel_size=3, padding=1)
+        self.enc2 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
+        self.enc3 = nn.Conv2d(64, 32, kernel_size=3, padding=1)
+        self.enc4 = nn.Conv2d(32, 16, kernel_size=3, padding=1)
+        self.pool = nn.MaxPool2d(2, 2, return_indices=True)
+
+        # Decoder
+        self.dec0 = nn.ConvTranspose2d(16, 32, kernel_size=2, stride=2)
+        self.dec1 = nn.ConvTranspose2d(32, 64, kernel_size=2, stride=2)
+        self.dec2 = nn.ConvTranspose2d(64, 128, kernel_size=2, stride=2)
+        self.dec3 = nn.ConvTranspose2d(128, 256, kernel_size=2, stride=2)
+        self.dec4 = nn.ConvTranspose2d(256, 3, kernel_size=2, stride=2)
+
+    def forward(self, x):
+        x, _ = self.pool(F.relu(self.enc0(x)))
+        x, _ = self.pool(F.relu(self.enc1(x)))
+        x, _ = self.pool(F.relu(self.enc2(x)))
+        x, _ = self.pool(F.relu(self.enc3(x)))
+        encoded = self.enc4(x)
+        x, _ = self.pool(F.relu(encoded))
+
+        x = F.relu(self.dec0(x))
+        x = F.relu(self.dec1(x))
+        x = F.relu(self.dec2(x))
+        x = F.relu(self.dec3(x))
+        x = torch.sigmoid(self.dec4(x))
+        return x, encoded

reset_dataset.py

@@ -0,0 +1,22 @@
+import os
+import shutil
+
+def reset_dataset():
+    dataset_path = "./dataset"
+    train_path = os.path.join(dataset_path, "train")
+    valid_path = os.path.join(dataset_path, "valid")
+
+    # Remove as pastas se existirem
+    if os.path.exists(train_path):
+        shutil.rmtree(train_path)
+    if os.path.exists(valid_path):
+        shutil.rmtree(valid_path)
+
+    # Recria as pastas
+    os.makedirs(train_path)
+    os.makedirs(valid_path)
+
+    print("O dataset foi resetado com sucesso!")
+
+if __name__ == "__main__":
+    reset_dataset()

run.py

@@ -12,8 +12,10 @@
     ConvolutionalAutoencoder,
     ConvolutionalVAE,
     DenoisingAutoencoder,
+    SparseAutoencoder,
     VariationalAutoencoder,
     DenoisingConvolutionalAutoencoder,
+    SparseConvolutionalAutoencoder
 )
 
 from settings import settings
@@ -26,12 +28,14 @@ def get_model_by_type(ae_type=None, input_dim=None, encoding_dim=None, device=No
     models = {
         'ae': lambda: Autoencoder(input_dim, encoding_dim),
         'dae': lambda: DenoisingAutoencoder(input_dim, encoding_dim),
+        'sparse': lambda: SparseAutoencoder(input_dim, encoding_dim),
         'vae': VariationalAutoencoder,
         'conv': ConvolutionalAutoencoder,
         'conv_dae': DenoisingConvolutionalAutoencoder,
         'conv_vae': ConvolutionalVAE,
+        'conv_sparse': SparseConvolutionalAutoencoder,
     }
-    
+
     if ae_type is None:
         return list(models.keys())
 
@@ -81,15 +85,15 @@ def main(load_trained_model, ae_type=None, num_epochs=5, test_mode=True):
     model = get_model_by_type(ae_type, input_dim, encoding_dim, device)
     optimizer = torch.optim.Adam(model.parameters())
 
-    start_epoch = 0
-    if os.path.exists(settings.PATH_SAVED_MODEL):
-        model, optimizer, start_epoch = load_checkpoint(
-            model, optimizer, settings.PATH_SAVED_MODEL, device
-        )
-        print(f"Loaded checkpoint and continuing training from epoch {start_epoch}.")
-
     try:
         if not load_trained_model:
+            start_epoch = 1
+            if os.path.exists(settings.PATH_SAVED_MODEL):
+                model, optimizer, start_epoch = load_checkpoint(
+                    model, optimizer, settings.PATH_SAVED_MODEL, device
+                )
+                print(f"Loaded checkpoint and continuing training from epoch {start_epoch}.")
+
             start_time = time.time()
 
             train_autoencoder(
@@ -100,7 +104,8 @@ def main(load_trained_model, ae_type=None, num_epochs=5, test_mode=True):
                 device=device,
                 start_epoch=start_epoch,
                 optimizer=optimizer,
-                save_checkpoint=save_checkpoint
+                save_checkpoint=save_checkpoint,
+                ae_type=ae_type
             )
 
             elapsed_time = utils.format_time(time.time() - start_time)
@@ -112,12 +117,12 @@ def main(load_trained_model, ae_type=None, num_epochs=5, test_mode=True):
 
     if not test_mode:
         valid_dataloader = get_dataloader(settings.VALID_DATA_PATH, batch_size, resolution)
-        avg_valid_loss = evaluate_autoencoder(model, valid_dataloader, device, ae_type)
+        avg_valid_loss = evaluate_autoencoder(model, valid_dataloader, device)
         print(f"\nAverage validation loss: {avg_valid_loss:.4f}\n")
 
         visualize_reconstructions(
             model, valid_dataloader, num_samples=10,
-            device=device, ae_type=ae_type, resolution=resolution
+            device=device, resolution=resolution
         )