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),
        )
Exemple #3
0
 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
     )