Views
No views yet
hidden_dim)1class ImageGenerationTAU(nn.Module, PyTorchModelHubMixin):
2 def __init__(self, hidden_dim):
3 super(ImageGenerationTAU, self).__init__()
4 self.encoder = nn.Sequential(
5 nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
6 nn.MaxPool2d(kernel_size=2, stride=2),
7 nn.ReLU(),
8 nn.BatchNorm2d(64),
9 nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1),
10 nn.MaxPool2d(kernel_size=2, stride=2),
11 nn.ReLU(),
12 nn.BatchNorm2d(32),
13 nn.Flatten(),
14 nn.Linear(32 * 7 * 7, hidden_dim),
15 )
16 self.decoder = nn.Sequential(
17 nn.Linear(hidden_dim, 32 * 7 * 7),
18 nn.ReLU(),
19 nn.Unflatten(1, (32, 7, 7)),
20 nn.ConvTranspose2d(32, 64, kernel_size=2, stride=2),
21 nn.ReLU(),
22 nn.BatchNorm2d(64),
23 nn.ConvTranspose2d(64, 1, kernel_size=2, stride=2),
24 nn.Sigmoid(),
25 )
26
27 def forward(self, x):
28 x = self.encoder(x)
29 x = self.decoder(x)
30 return x