in models.py [0:0]
def __init__(self, n_pixels, n_channels, z_dim, distribution="gaussian"):
super().__init__()
self.z_dim = z_dim
self.n_channels = n_channels
self.distribution = distribution
self.encoder = nn.Sequential(
nn.Conv2d(n_channels, n_pixels, kernel_size=4, stride=2),
nn.ReLU(),
nn.Conv2d(n_pixels, n_pixels, kernel_size=4, stride=2),
nn.ReLU(),
nn.Conv2d(n_pixels, n_pixels, kernel_size=4, stride=2),
nn.ReLU(),
nn.Conv2d(n_pixels, 256, kernel_size=1, stride=1),
nn.ReLU(),
Lambda(lambda x: x.view(x.size(0), -1)),
nn.Linear(256, 2 * z_dim),
)
self.decoder = nn.Sequential(
nn.Linear(z_dim, 256),
nn.ReLU(),
Lambda(lambda x: x.view(-1, 256, 1, 1)),
nn.ConvTranspose2d(256, 64, 4),
nn.ReLU(),
nn.ConvTranspose2d(64, 64, 4, 2, 1),
nn.ReLU(),
nn.ConvTranspose2d(64, n_pixels, 4, 2, 1),
nn.ReLU(),
nn.ConvTranspose2d(n_pixels, n_channels, 4, 2, 1),
Lambda(lambda x: x.view(x.size(0), -1)),
nn.ReLU(),
nn.Linear(32 * 32, n_pixels * n_pixels),
Lambda(lambda x: x.view(x.size(0), 1, n_pixels, n_pixels)),
nn.Sigmoid(),
)