Merge pull request #13 from renan-siqueira/feature/UpgradeArchitecture

Upgrade architecture

Author: renan-siqueira

README.md

@@ -2,7 +2,7 @@
 
 A simple implementation of an autoencoder using PyTorch. 
 
-This project aims to provide a foundational structure to understand, train, and evaluate autoencoders on 64x64 images.
+This project aims to provide a basic framework for understanding, training and evaluating autoencoders on any image size.
 
 ## Features
 
@@ -25,25 +25,15 @@ This project aims to provide a foundational structure to understand, train, and
 
 ### Installation
 
-1. Clone the repository:
-
-```bash
-git clone https://github.com/renan-siqueira/autoencoder-project.git
-```
-2. Navigate to the project directory and install the required libraries:
-
-```bash
-cd autoencoder-project
-pip install -r requirements.txt
-```
+1. Clone the repository.
+2. Navigate to the project directory and install the required libraries.
 
 ## Usage
 
-1. Modify settings/settings.py to point to your training and validation dataset.
-2. To train the autoencoder, simply run:
+1. Modify the `settings/settings.py` file to point to your training and validation dataset.
+2. Modify the `json/params.json` file to reflext your training preferences. 
+3. To train the autoencoder, simply run:
 
 ```bash
 python run.py
 ```
-
-By default, this will train a new model. If you wish to use a pre-trained model, modify the `main` method in `run.py`.

json/params.json

@@ -4,6 +4,6 @@
     "encoding_dim": 16,
     "num_epochs": 1000,
     "learning_rate": 0.001,
-    "ae_type": "ae",
+    "ae_type": "conv",
     "save_checkpoint": null
 }

models/convolutional_autoencoder.py

@@ -8,22 +8,30 @@ def __init__(self):
         super(ConvolutionalAutoencoder, self).__init__()
 
         # Encoder
-        self.enc1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
-        self.enc2 = nn.Conv2d(64, 32, kernel_size=3, padding=1)
-        self.enc3 = nn.Conv2d(32, 16, kernel_size=3, padding=1)
+        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.dec1 = nn.ConvTranspose2d(16, 32, kernel_size=2, stride=2)
-        self.dec2 = nn.ConvTranspose2d(32, 64, kernel_size=2, stride=2)
-        self.dec3 = nn.ConvTranspose2d(64, 3, kernel_size=2, stride=2)
+        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)))
+        x, _ = self.pool(F.relu(self.enc4(x)))
 
+        x = F.relu(self.dec0(x))
         x = F.relu(self.dec1(x))
         x = F.relu(self.dec2(x))
-        x = torch.sigmoid(self.dec3(x))
+        x = F.relu(self.dec3(x))
+        x = torch.sigmoid(self.dec4(x))
         return x