backup

src/models/__init__.py

@@ -1,8 +1,10 @@
 from .Model import Model
+from .autoencoder import AutoEncoder
 from .cnn import CNNPredictor
 from .transformer import ByteTransformer
 
 model_called: dict[str, type[Model]] = {
     'cnn': CNNPredictor,
-    'transformer': ByteTransformer
+    'transformer': ByteTransformer,
+    'autoencoder': AutoEncoder
 }

src/models/autoencoder/__init__.py

@@ -0,0 +1 @@
+from .autoencoder import AutoEncoder

src/models/autoencoder/autoencoder.py

@@ -1,6 +1,8 @@
 import torch
 import torch.nn as nn
 
+from src.models import Model
+
 
 class Encoder(nn.Module):
     def __init__(self, input_size, hidden_size, latent_dim):
@@ -23,21 +25,21 @@ class Decoder(nn.Module):
     def __init__(self, input_size, hidden_size, output_size):
         super(Decoder, self).__init__()
         super._decoder = nn.Sequential(*[
-            nn.Linear(input_size),
+            nn.Linear(input_size, 2 * hidden_size),
             nn.ReLU(),
-            nn.BatchNorm1d(input_size),
-            nn.ConvTranspose1d(input_size, 2 * hidden_size, kernel_size=3, stride=2, padding=1, output_padding=1),
             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.ReLU(),
-            nn.ConvTranspose1d(2 * hidden_size, output_size, kernel_size=3, padding=1),
+            nn.ConvTranspose1d(hidden_size, output_size, kernel_size=3, padding=1),
         ])
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        pass
+        return self._decoder(x)
 
-class AutoEncoder(nn.Module):
+class AutoEncoder(Model):
     def __init__(self, input_size, hidden_size, latent_dim):
-        super(AutoEncoder, self).__init__()
+        super().__init__(loss_function = nn.CrossEntropyLoss())
 
         self.encoder = Encoder(input_size, hidden_size, latent_dim)
         self.decoder = Decoder(latent_dim, hidden_size, input_size)