Merge pull request #7 from renan-siqueira/feature/Improvements

Checkpoint Feature

Author: renan-siqueira

run.py

@@ -3,36 +3,37 @@
 
 from models.autoencoder import Autoencoder
 from utils.dataloader import get_dataloader
-from utils.trainer import train_autoencoder, visualize_reconstructions, save_model, load_model, evaluate_autoencoder
+from utils.trainer import train_autoencoder, visualize_reconstructions, save_checkpoint, load_checkpoint, evaluate_autoencoder
 from settings import settings
 
 
 def main(load_trained_model):
     BATCH_SIZE = 32
     INPUT_DIM = 3 * 64 * 64
-    ENCODING_DIM = 12
-    NUM_EPOCHS = 1000
+    ENCODING_DIM = 64
+    NUM_EPOCHS = 200
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
     dataloader = get_dataloader(settings.DATA_PATH, BATCH_SIZE)
-    model = Autoencoder(INPUT_DIM, ENCODING_DIM).to(device)
 
-    if load_trained_model:
-        trained_model = load_model(model, settings.PATH_SAVED_MODEL, device=device)
-    else:
-        trained_model = train_autoencoder(model, dataloader, NUM_EPOCHS, device=device)
+    model = Autoencoder(INPUT_DIM, ENCODING_DIM).to(device)
+    optimizer = torch.optim.Adam(model.parameters())
 
-    valid_dataloader = get_dataloader(settings.VALID_DATA_PATH, BATCH_SIZE)
+    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}.")
 
-    save_path = os.path.join('./', settings.PATH_SAVED_MODEL)
-    save_model(trained_model, save_path)
-    print(f"Model saved to {save_path}")
+    if not load_trained_model:
+        for epoch in range(start_epoch, NUM_EPOCHS):
+            train_autoencoder(model, dataloader, device=device)
+            print(f"Epoch [{epoch + 1}/{NUM_EPOCHS}] complete!")
+            save_checkpoint(model, optimizer, epoch, settings.PATH_SAVED_MODEL)
 
-    avg_valid_loss = evaluate_autoencoder(trained_model, valid_dataloader, device)
+    valid_dataloader = get_dataloader(settings.VALID_DATA_PATH, BATCH_SIZE)
+    avg_valid_loss = evaluate_autoencoder(model, valid_dataloader, device)
     print(f"Average validation loss: {avg_valid_loss:.4f}")
-
-    visualize_reconstructions(trained_model, valid_dataloader, num_samples=10, device=device)
+    visualize_reconstructions(model, valid_dataloader, num_samples=10, device=device)
 
 
 if __name__ == "__main__":

utils/trainer.py

@@ -2,17 +2,16 @@
 import torch
 import torch.optim as optim
 import torch.nn as nn
-from torchvision import transforms
 from torchvision.utils import save_image, make_grid
 import matplotlib.pyplot as plt
-from PIL import Image
 
 
-def train_autoencoder(model, dataloader, num_epochs=5, learning_rate=0.001, device='cpu'):
+def train_autoencoder(model, dataloader, num_epochs=5, learning_rate=0.001, device='cpu', start_epoch=0, optimizer=None):
     criterion = nn.MSELoss()
-    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+    if optimizer is None:
+        optimizer = optim.Adam(model.parameters(), lr=learning_rate)
 
-    for epoch in range(num_epochs):
+    for epoch in range(start_epoch, num_epochs):
         for data in dataloader:
             img = data.to(device)
             img = img.view(img.size(0), -1)
@@ -24,6 +23,7 @@ def train_autoencoder(model, dataloader, num_epochs=5, learning_rate=0.001, devi
             optimizer.step()
 
         print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')
+        save_checkpoint(model, optimizer, epoch, './autoencoder_checkpoint.pth')
 
     return model
 
@@ -38,11 +38,9 @@ def visualize_reconstructions(model, dataloader, num_samples=10, device='cpu', s
     samples = samples.view(-1, 3, 64, 64)
     reconstructions = reconstructions.view(-1, 3, 64, 64)
 
-    # Combine as amostras e reconstruções em uma única grade
     combined = torch.cat([samples, reconstructions], dim=0)
     grid_img = make_grid(combined, nrow=num_samples)
 
-    # Visualização usando Matplotlib
     plt.imshow(grid_img.permute(1, 2, 0).cpu().detach().numpy())
     plt.axis('off')
     plt.show()
@@ -62,6 +60,23 @@ def load_model(model, path, device):
     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
+
+
 def evaluate_autoencoder(model, dataloader, device):
     model.eval()
     total_loss = 0