changes to training + added autoencoder

src/models/autoencoder/autoencoder.py

@@ -5,15 +5,16 @@ from src.models import Model
 
 
 class Encoder(nn.Module):
-    def __init__(self, input_size, hidden_size, latent_dim):
+    def __init__(self, data_length, channel_count, latent_dim):
         super(Encoder, self).__init__()
         self._encoder = nn.Sequential(*[
-            nn.Conv1d(input_size, hidden_size, kernel_size=3, padding=1),
-            nn.BatchNorm1d(hidden_size),
+            nn.Conv1d(1, channel_count, kernel_size=3, padding=1), # (hidden_size, L)
+            nn.BatchNorm1d(channel_count),
             nn.ReLU(),
-            nn.Conv1d(hidden_size, 2 * hidden_size, stride=2, kernel_size=3, padding=1),
-            nn.BatchNorm1d(2 * hidden_size),
-            nn.Linear(2 * hidden_size, latent_dim),
+            nn.Conv1d(channel_count, 2 * channel_count, stride=2, kernel_size=3, padding=1), # (2 * hidden_size, L / 2)
+            nn.BatchNorm1d(2 * channel_count),
+            nn.Flatten(), # 2 * hidden_size * L / 2
+            nn.Linear(2 * channel_count * data_length // 2, latent_dim),
             nn.ReLU()
         ])
 
@@ -22,27 +23,28 @@ class Encoder(nn.Module):
 
 
 class Decoder(nn.Module):
-    def __init__(self, input_size, hidden_size, output_size):
+    def __init__(self, latent_dim, channel_count, data_length):
         super(Decoder, self).__init__()
-        super._decoder = nn.Sequential(*[
-            nn.Linear(input_size, 2 * hidden_size),
+        self._decoder = nn.Sequential(*[
+            nn.Linear(latent_dim, 2 * channel_count * data_length // 2),
             nn.ReLU(),
-            nn.BatchNorm1d(2 * hidden_size),
-            nn.ConvTranspose1d(2 * hidden_size, hidden_size, kernel_size=3, stride=2, padding=1, output_padding=1),
-            nn.BatchNorm1d(hidden_size),
+            nn.Unflatten(1, (2 * channel_count, data_length // 2)),
+            nn.BatchNorm1d(2 * channel_count),
+            nn.ConvTranspose1d(2 * channel_count, channel_count, kernel_size=3, stride=2, padding=1, output_padding=1),
+            nn.BatchNorm1d(channel_count),
             nn.ReLU(),
-            nn.ConvTranspose1d(hidden_size, output_size, kernel_size=3, padding=1),
+            nn.ConvTranspose1d(channel_count, 1, kernel_size=3, padding=1),
         ])
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self._decoder(x)
 
 class AutoEncoder(Model):
-    def __init__(self, input_size, hidden_size, latent_dim):
-        super().__init__(loss_function = nn.CrossEntropyLoss())
+    def __init__(self, input_size, channel_count, latent_dim):
+        super().__init__(loss_function = nn.MSELoss())
 
-        self.encoder = Encoder(input_size, hidden_size, latent_dim)
-        self.decoder = Decoder(latent_dim, hidden_size, input_size)
+        self.encoder = Encoder(input_size, channel_count, latent_dim)
+        self.decoder = Decoder(latent_dim, channel_count, input_size)
 
     def encode(self, x: torch.Tensor) -> torch.Tensor:
         return self.encoder(x)
@@ -50,5 +52,11 @@ class AutoEncoder(Model):
     def decode(self, x: torch.Tensor) -> torch.Tensor:
         return self.decoder(x)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.decode(self.encode(x))
+    def forward(self, x: torch.LongTensor) -> torch.Tensor:
+        x = x.float() / 255.0 # convert to floats
+        x = x.unsqueeze(1) # add channel dimension --> (B, 1, L)
+
+        encoded = self.encoder(x)
+        decoded = self.decoder(encoded)
+
+        return decoded