Create model.py

model.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pytorch_lightning as pl
+
+# ----------------------------------------------------
+# A helper block for the Residual Connection
+# ----------------------------------------------------
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+
+        # Skip connection for differing channels/stride
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_channels != out_channels:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channels)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+# ----------------------------------------------------
+# The Main Residual Autoencoder Model
+# ----------------------------------------------------
+class ResidualConvAutoencoder(pl.LightningModule):
+    def __init__(self, latent_dim=512, dropout_rate=0.2):
+        super().__init__()
+        self.latent_dim = latent_dim
+        
+        # --- Encoder ---
+        self.encoder = nn.Sequential(
+            nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1), # 32x32 -> 32x32
+            ResidualBlock(64, 128, stride=2), # 32x32 -> 16x16
+            ResidualBlock(128, 256, stride=2), # 16x16 -> 8x8
+            ResidualBlock(256, 512, stride=2), # 8x8 -> 4x4
+            nn.Flatten(),
+            nn.Linear(512 * 4 * 4, self.latent_dim),
+            nn.Dropout(dropout_rate)
+        )
+        
+        # --- Decoder ---
+        self.decoder = nn.Sequential(
+            nn.Linear(self.latent_dim, 512 * 4 * 4),
+            nn.Unflatten(1, (512, 4, 4)),
+            ResidualBlock(512, 256),
+            nn.Upsample(scale_factor=2, mode='nearest'), # 4x4 -> 8x8
+            ResidualBlock(256, 128),
+            nn.Upsample(scale_factor=2, mode='nearest'), # 8x8 -> 16x16
+            ResidualBlock(128, 64),
+            nn.Upsample(scale_factor=2, mode='nearest'), # 16x16 -> 32x32
+            nn.Conv2d(64, 3, kernel_size=3, stride=1, padding=1),
+            nn.Sigmoid() # Output pixel values between 0 and 1
+        )
+        
+    def forward(self, x):
+        z = self.encoder(x)
+        recon = self.decoder(z)
+        return recon
+    
+    # Placeholder training step (not needed for deployment file, but required for class completeness)
+    def training_step(self, batch, batch_idx):
+        return torch.tensor(0.0) 
+
+    # Placeholder configure_optimizers (not needed for deployment file, but required for class completeness)
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters())