2nd commit ,# Plot Training Loss vs. Epochs

Standard autoencoder.py → Standard Autoencoder _CNN_ENCODER-DECODER.py

@@ -22,8 +22,7 @@ transform = transforms.Compose([
 # Hyperparameters
 batch_size = 100
 IMAGE_SIZE = 64
-# Define the path to the dataset
-IMAGE_DIRECTORY = '/home/mhossain/PycharmProjects/python_Project_pendulum/dataset'
+IMAGE_DIRECTORY = '/home/mhossain/PycharmProjects/python_Project_pendulum/dataset' # Define the path to the dataset
 
 
 #Define your dataset path and create DataLoader for training, validation and testing
@@ -37,7 +36,7 @@ val_dataloader = DataLoader(Subset(dataset, range(16001, 18000)), batch_size=bat
 test_dataloader = DataLoader(Subset(dataset, range(18001, len(dataset))), batch_size=batch_size, num_workers=4, shuffle=False)
 
 
-#Verify Dataset Size
+#Verify number of images in the dataset
 print(f"Total number of images in the dataset: {len(dataset)}")
 
 
@@ -49,13 +48,13 @@ class AutoEncoder(pl.LightningModule):
         # Encoder # Input size = 1 x 64 x 64
 
         self.encoder = nn.Sequential(
-            nn.Conv2d(1, 32, kernel_size=4, stride=2, padding=1),
+            nn.Conv2d(1, 32, kernel_size=4, stride=2, padding=1), # Output size = 32 x 32 x 32
             nn.ReLU(),
-            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),
+            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1), # Output size = 64 x 16 x 16
             nn.ReLU(),
-            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
+            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1), # Output size = 128 x 8 x 8
             nn.ReLU(),
-            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
+            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),  # Output size = 256 x 4 x 4
             nn.ReLU(),
             nn.Flatten(),
             nn.Linear(256 * 4 * 4, 8)  # Latent dimension size = 8
@@ -66,62 +65,99 @@ class AutoEncoder(pl.LightningModule):
             nn.Linear(8, 256 * 4 * 4),  # Latent dimension size = 8
             nn.ReLU(),
             nn.Unflatten(1, (256, 4, 4)),
-            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),
+            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1), # Output size = 128 x 8 x 8
             nn.ReLU(),
-            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
+            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1), # Output size = 64 x 16 x 16
             nn.ReLU(),
-            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),
+            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1), # Output size = 32 x 32 x 32
             nn.ReLU(),
-            nn.ConvTranspose2d(32, 1, kernel_size=4, stride=2, padding=1),
+            nn.ConvTranspose2d(32, 1, kernel_size=4, stride=2, padding=1), # Output size = 1 x 64 x 64
             nn.Sigmoid() # Output values between 0 and 1
         )
 
 
     def forward(self, x): # Define the forward pass
-        z = self.encoder(x)
-        recon_x = self.decoder(z)
+        z = self.encoder(x) # Encode the input image to a latent vector
+        recon_x = self.decoder(z) # Decode the latent vector
         return recon_x
 
     def loss_function(self, recon_x, x): # Define the loss function
         # Reconstruction loss
         BCE = F.binary_cross_entropy(recon_x, x, reduction='mean')
+        # Calculate the binary cross entropy loss between the reconstructed image and the original image
         return BCE
 
     def training_step(self, batch, batch_idx):# Define the training step
-        x, _ = batch
-        recon_x = self(x)
-        loss = self.loss_function(recon_x, x)
-        self.log('train_loss', loss)
+        x, _ = batch # Get the input image and its corresponding label
+        recon_x = self(x) # Reconstruct the input image
+        loss = self.loss_function(recon_x, x) # Calculate the reconstruction loss
+        self.log('train_loss', loss) # Log the training loss
         return loss
 
     def validation_step(self, batch, batch_idx): # Define the validation step
-        x, _ = batch
-        recon_x = self(x)
-        loss = self.loss_function(recon_x, x)
-        self.log('val_loss', loss)
+        x, _ = batch # Get the input image and its corresponding label
+        recon_x = self(x) # Reconstruct the input image
+        loss = self.loss_function(recon_x, x) # Calculate the reconstruction loss
+        self.log('val_loss', loss) # Log the validation loss
 
     def configure_optimizers(self):  # Define the optimizer
-        return torch.optim.Adam(self.parameters(), lr=0.002)
+        return torch.optim.Adam(self.parameters(), lr=0.002)  # Use the Adam optimizer with a learning rate of 0.002
 
 
-# Initialize the AutoEncoder model
 
-model = AutoEncoder()
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+ # Callback to plot the training loss vs. epochs
+class TrainingLossPlotCallback(pl.Callback):
+    def __init__(self, max_epochs):
+        super().__init__()
+        self.max_epochs = max_epochs
+        self.train_losses = []
+
+    def on_epoch_end(self, trainer, pl_module):
+        # Record the training loss for the current epoch
+        train_loss = trainer.callback_metrics.get('train_loss')
+        if train_loss is not None:
+            self.train_losses.append(train_loss.item())
+
+    def on_train_end(self, trainer, pl_module):
+        # Plot Training Loss vs. Epochs
+        plt.figure(figsize=(10, 6))
+        plt.plot(range(1, len(self.train_losses) + 1), self.train_losses, label='Training Loss', color='orange')
+        plt.xlabel('Epochs')
+        plt.ylabel('Loss')
+        plt.title('Training Loss vs Epochs')
+        plt.xlim(1, self.max_epochs)  # Set x-axis limit to max_epochs
+        plt.legend()
+        plt.savefig('training_loss_vs_epochs.png')
+        plt.close()
+        plt.clf()
 
-from torchsummary import summary
-summary(model, input_size=(1, 64, 64))
 
+# Initialize the AutoEncoder model
+model = AutoEncoder()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # Use the GPU if available
+model.to(device) # Transfer the model to the GPU if available
+
+from torchsummary import summary # Print the model summary to the console
+summary(model, input_size=(1, 64, 64))  # 1 x 64 x 64 is the input image dimension
+
+
+max_epochs = 100 # Define the maximum number of epochs
+# Integrating the Callback with the Trainer
+training_loss_plot_callback = TrainingLossPlotCallback(max_epochs)
+# Initialize the trainer
+trainer = pl.Trainer(
+    gpus=1,
+    max_epochs=max_epochs,
+    callbacks=[training_loss_plot_callback]
+)
+# Train the model
+trainer.fit(model, train_dataloader, val_dataloader)
 
-#PyTorch Lightning trainer, specifies GPU usage and number of epochs.
 
-trainer = pl.Trainer(gpus=1, max_epochs=50)
-trainer.fit(model, train_dataloader, val_dataloader)
 
 
 # Function to save a batch of images with labels to a file
-def save_images(original_data, reconstructed_data, filename='comparison.png', num_images=2):
+def save_images(original_data, reconstructed_data, filename="reconstructed_images.png", num_images=5):
     plt.figure(figsize=(12, 4))
 
     for i in range(num_images):
@@ -142,6 +178,7 @@ def save_images(original_data, reconstructed_data, filename='comparison.png', nu
         # Save the entire figure
     plt.savefig(filename)
     plt.close()
+    plt.clf()
 
 
 # Randomly select a batch of images
@@ -153,4 +190,4 @@ with torch.no_grad():
     loss = model.loss_function(reconstructed_data, data)  # Calculate the reconstruction loss
     print(f"Reconstruction Loss: {loss.item()}")  # Print the reconstruction loss
 
-    save_images(data, reconstructed_data, filename="reconstructed_images.png") # Save the reconstructed images to a file
\ No newline at end of file
+    save_images(data, reconstructed_data, filename="reconstructed_images.png") # Save the reconstructed images to a file