def __init__(self, z_dim=10, nc=3, target_size=(128, 128)):
super(VanillaVAEDecoder, self).__init__()
self.target_size = target_size
self.decoder = nn.Sequential(
nn.Linear(z_dim, 256), # B, 256
nn.ReLU(True),
nn.Linear(256, 256), # B, 256
nn.ReLU(True),
nn.Linear(256, 32 * 8 * 8), # B, 2048
nn.ReLU(True),
View((-1, 32, 8, 8)), # B, 32, 8, 8
nn.ConvTranspose2d(32, 32, 4, 2, 1), # B, 32, 8, 8
nn.ReLU(True),
nn.ConvTranspose2d(32, 32, 4, 2, 1), # B, 32, 16, 16
nn.ReLU(True),
nn.ConvTranspose2d(32, 32, 4, 2, 1), # B, 32, 32, 32
nn.ReLU(True),
nn.ConvTranspose2d(32, nc, 4, 2, 1), # B, nc, 64, 64
nn.Tanh(),
View(self.target_size),
)
def __init__(self, z_dim, nc, target_size):
super(WassersteinADecoder, self).__init__()
self.z_dim = z_dim
self.nc = nc
self.target_size = target_size
self.decoder = nn.Sequential(
nn.Linear(z_dim, 256), # B, 256
nn.ReLU(True),
nn.Linear(256, 256), # B, 256
nn.ReLU(True),
nn.Linear(256, 32 * 8 * 8), # B, 2048
nn.ReLU(True),
View((-1, 32, 8, 8)), # B, 32, 8, 8
nn.ConvTranspose2d(32, 32, 4, 2, 1), # B, 32, 8, 8
nn.ReLU(True),
nn.ConvTranspose2d(32, 32, 4, 2, 1), # B, 32, 16, 16
nn.ReLU(True),
nn.ConvTranspose2d(32, 32, 4, 2, 1), # B, 32, 32, 32
nn.ReLU(True),
nn.ConvTranspose2d(32, nc, 4, 2, 1), # B, nc, 64, 64
nn.Tanh(),
View(self.target_size),
)
def __init__(self, z_dim, nc):
super(AAEEncoder, self).__init__()
self.encoder = nn.Sequential(
nn.Conv2d(nc, 32, 4, 2, 1), # B, 32, 32, 32
nn.ReLU(True),
nn.Conv2d(32, 32, 4, 2, 1), # B, 32, 16, 16
nn.ReLU(True),
nn.Conv2d(32, 32, 4, 2, 1), # B, 32, 8, 8
nn.ReLU(True),
nn.Conv2d(32, 32, 4, 2, 1), # B, 32, 8, 8
nn.ReLU(True),
View((-1, 32 * 8 * 8)), # B, 2048
nn.Linear(32 * 8 * 8, 512), # B, 512
nn.ReLU(True),
nn.Linear(512, 256), # B, 256
nn.ReLU(True),
nn.Linear(256, z_dim), # B, z_dim*2
)