| 24 | # %% model class |
| 25 | LATENT_DIMS = 128 |
| 26 | class Encoder(nn.Module): |
| 27 | def __init__(self) -> None: |
| 28 | super().__init__() |
| 29 | self.conv1 = nn.Conv2d(1, 6, 3) # out: 6, 62, 62 |
| 30 | self.conv2 = nn.Conv2d(6, 16, 3) # out: 16, 60, 60 |
| 31 | self.relu = nn.ReLU() |
| 32 | self.flatten = nn.Flatten() # out: 16*60*60 = 57600 |
| 33 | self.fc = nn.Linear(16*60*60, LATENT_DIMS) |
| 34 | |
| 35 | |
| 36 | def forward(self, x): |
| 37 | x = self.conv1(x) |
| 38 | x = self.relu(x) |
| 39 | x = self.conv2(x) |
| 40 | x = self.relu(x) |
| 41 | x = self.flatten(x) |
| 42 | x = self.fc(x) |
| 43 | return x |
| 44 | |
| 45 | class Decoder(nn.Module): |
| 46 | def __init__(self) -> None: |