feat: upgrade process and adapted utils

Author: renan-siqueira

json/params.json

@@ -1,8 +1,9 @@
 {
     "batch_size": 32,
     "resolution": 64,
-    "encoding_dim": 128,
-    "num_epochs": 500,
+    "encoding_dim": 16,
+    "num_epochs": 1000,
     "learning_rate": 0.001,
-    "ae_type": "conv_vae"
+    "ae_type": "ae",
+    "save_checkpoint": null
 }

run.py

@@ -1,12 +1,25 @@
 import os
 import json
+import random
+import time
 
+import numpy as np
 import torch
 
 from models import Autoencoder, ConvolutionalAutoencoder, ConvolutionalVAE, VariationalAutoencoder
+from settings import settings
 from utils.dataloader import get_dataloader
 from utils.trainer import train_autoencoder, visualize_reconstructions, load_checkpoint, evaluate_autoencoder
-from settings import settings
+from utils import utils
+
+
+def set_seed(seed=42):
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
 
 
 def load_params(path):
@@ -16,6 +29,7 @@ def load_params(path):
 
 
 def main(load_trained_model):
+    set_seed(1)
     params = load_params(settings.PATH_PARAMS_JSON)
 
     batch_size = params["batch_size"]
@@ -24,12 +38,13 @@ def main(load_trained_model):
     num_epochs = params["num_epochs"]
     learning_rate = params.get("learning_rate", 0.001)
     ae_type = params["ae_type"]
+    save_checkpoint = params["save_checkpoint"]
 
     # Calculate input_dim based on resolution
     input_dim = 3 * resolution * resolution
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    dataloader = get_dataloader(settings.DATA_PATH, batch_size)
+    dataloader = get_dataloader(settings.DATA_PATH, batch_size, resolution)
 
     if ae_type == 'ae':
         model = Autoencoder(input_dim, encoding_dim).to(device)
@@ -51,25 +66,35 @@ def main(load_trained_model):
         )
         print(f"Loaded checkpoint and continuing training from epoch {start_epoch}.")
 
-    if not load_trained_model:
-        train_autoencoder(
-            model,
-            dataloader,
-            num_epochs=num_epochs,
-            learning_rate=learning_rate,
-            device=device,
-            start_epoch=start_epoch,
-            optimizer=optimizer,
-            ae_type=ae_type
-        )
-        print(f"Training complete up to epoch {num_epochs}!")
-
-    valid_dataloader = get_dataloader(settings.VALID_DATA_PATH, batch_size)
+    try:
+        if not load_trained_model:
+            start_time = time.time()
+
+            train_autoencoder(
+                model,
+                dataloader,
+                num_epochs=num_epochs,
+                learning_rate=learning_rate,
+                device=device,
+                start_epoch=start_epoch,
+                optimizer=optimizer,
+                ae_type=ae_type,
+                save_checkpoint=save_checkpoint
+            )
+
+            elapsed_time = utils.format_time(time.time() - start_time)
+            print(f"\nTraining took {elapsed_time}")
+            print(f"Training complete up to epoch {num_epochs}!")
+    
+    except KeyboardInterrupt:
+        print("\nTraining interrupted by user.")
+
+    valid_dataloader = get_dataloader(settings.VALID_DATA_PATH, batch_size, resolution)
     avg_valid_loss = evaluate_autoencoder(model, valid_dataloader, device, ae_type)
-    print(f"Average validation loss: {avg_valid_loss:.4f}")
+    print(f"\nAverage validation loss: {avg_valid_loss:.4f}\n")
 
     visualize_reconstructions(
-        model, valid_dataloader, num_samples=10, device=device, ae_type=ae_type
+        model, valid_dataloader, num_samples=10, device=device, ae_type=ae_type, resolution=resolution
     )
 
 

utils/dataloader.py

@@ -1,13 +1,11 @@
 import os
-import torch
-from torchvision import datasets, transforms
 from torchvision.transforms import ToTensor, Resize, Compose
 from torch.utils.data import DataLoader, Dataset
 from PIL import Image
 
 
-def get_dataloader(data_path, batch_size):
-    dataset = CustomDataset(data_path)
+def get_dataloader(data_path, batch_size, resolution):
+    dataset = CustomDataset(data_path, resolution)
 
     dataloader = DataLoader(
         dataset,
@@ -19,12 +17,12 @@ def get_dataloader(data_path, batch_size):
 
 
 class CustomDataset(Dataset):
-    def __init__(self, data_path):
+    def __init__(self, data_path, resolution):
         self.data_path = data_path
         self.image_files = os.listdir(data_path)
 
         self.transforms = Compose([
-            Resize((64, 64)),
+            Resize((resolution, resolution)),
             ToTensor()
         ])
 

utils/trainer.py

@@ -7,7 +7,37 @@
 import matplotlib.pyplot as plt
 
 
-def train_autoencoder(model, dataloader, num_epochs, learning_rate, device, start_epoch, optimizer, ae_type):
+def save_model(model, path):
+    torch.save(model.state_dict(), path)
+
+
+def load_model(model, path, device):
+    model.load_state_dict(torch.load(path, map_location=device))
+    model.eval()
+    return model
+
+
+def save_checkpoint_file(model, optimizer, epoch, path):
+    checkpoint = {
+        'epoch': epoch,
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': optimizer.state_dict(),
+    }
+    torch.save(checkpoint, path)
+
+
+def load_checkpoint(model, optimizer, path, device):
+    checkpoint = torch.load(path, map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    epoch = checkpoint['epoch']
+    return model, optimizer, epoch + 1
+
+
+def train_autoencoder(
+        model, dataloader, num_epochs, learning_rate, device,
+        start_epoch, optimizer, ae_type, save_checkpoint
+    ):
     criterion = nn.MSELoss()
     if optimizer is None:
         optimizer = optim.Adam(model.parameters(), lr=learning_rate)
@@ -31,7 +61,8 @@ def train_autoencoder(model, dataloader, num_epochs, learning_rate, device, star
             optimizer.step()
 
         print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')
-        save_checkpoint(model, optimizer, epoch, './autoencoder_checkpoint.pth')
+        if save_checkpoint:
+            save_checkpoint_file(model, optimizer, epoch, './autoencoder_checkpoint.pth')
 
     return model
 
@@ -63,7 +94,7 @@ def evaluate_autoencoder(model, dataloader, device, ae_type):
     return total_loss / len(dataloader)
 
 
-def visualize_reconstructions(model, dataloader, num_samples=10, device='cpu', save_path="./samples", ae_type='ae'):
+def visualize_reconstructions(model, dataloader, num_samples=10, device='cpu', save_path="./samples", ae_type='ae', resolution=64):
     model.eval()
     samples = next(iter(dataloader))
     samples = samples[:num_samples].to(device)
@@ -76,8 +107,8 @@ def visualize_reconstructions(model, dataloader, num_samples=10, device='cpu', s
     else:
         reconstructions = model(samples)
 
-    samples = samples.view(-1, 3, 64, 64)
-    reconstructions = reconstructions.view(-1, 3, 64, 64)
+    samples = samples.view(-1, 3, resolution, resolution)
+    reconstructions = reconstructions.view(-1, 3, resolution, resolution)
 
     combined = torch.cat([samples, reconstructions], dim=0)
     grid_img = make_grid(combined, nrow=num_samples)
@@ -89,30 +120,3 @@ def visualize_reconstructions(model, dataloader, num_samples=10, device='cpu', s
     if not os.path.exists(save_path):
         os.makedirs(save_path)
     save_image(grid_img, os.path.join(save_path, 'combined_samples.png'))
-
-
-def save_model(model, path):
-    torch.save(model.state_dict(), path)
-
-
-def load_model(model, path, device):
-    model.load_state_dict(torch.load(path, map_location=device))
-    model.eval()
-    return model
-
-
-def save_checkpoint(model, optimizer, epoch, path):
-    checkpoint = {
-        'epoch': epoch,
-        'model_state_dict': model.state_dict(),
-        'optimizer_state_dict': optimizer.state_dict(),
-    }
-    torch.save(checkpoint, path)
-
-
-def load_checkpoint(model, optimizer, path, device):
-    checkpoint = torch.load(path, map_location=device)
-    model.load_state_dict(checkpoint['model_state_dict'])
-    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
-    epoch = checkpoint['epoch']
-    return model, optimizer, epoch + 1

utils/utils.py

@@ -0,0 +1,13 @@
+def format_time(elapsed_time):
+    if elapsed_time < 60:
+        return f"{elapsed_time:.2f} seconds"
+    elif elapsed_time < 3600:
+        minutes = elapsed_time // 60
+        seconds = elapsed_time % 60
+        return f"{minutes:.0f} minutes, {seconds:.2f} seconds"
+    else:
+        hours = elapsed_time // 3600
+        remainder = elapsed_time % 3600
+        minutes = remainder // 60
+        seconds = remainder % 60
+        return f"{hours:.0f} hours, {minutes:.0f} minutes, {seconds:.2f} seconds"