models.py

import torch
import torch.nn as nn
import torch.nn.functional as F

class VAE(nn.Module):
    def __init__(self, zdim = 60):
        super(VAE, self).__init__()

        #encoder
        self.conv1 = nn.Conv2d(3, 32, kernel_size=4, stride=2)
        self.conv2 = nn.Conv2d(32, 32, kernel_size=4, stride=2)
        self.conv3 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
        self.conv4 = nn.Conv2d(64, 64, kernel_size=4, stride=2)
        self.conv5 = nn.Conv2d(64, 128, kernel_size=4, stride=2)
        self.act = nn.ReLU()

        self.fc1 = nn.Linear(128 * 2 * 2, 256)
        self.fc2_mu = nn.Linear(256, zdim)
        self.fc2_logvar = nn.Linear(256, zdim)

        self.fc3 = nn.Linear(60, 128 * 4 * 4)

        #decoder
        self.deconv1 = nn.ConvTranspose2d(128, 128, kernel_size=3, stride=2, padding=1, output_padding=1)
        self.deconv2 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1)
        self.deconv3 = nn.ConvTranspose2d(64, 64, kernel_size=3, stride=2, padding=1, output_padding=1)
        self.deconv4 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1)
        self.deconv5 = nn.ConvTranspose2d(32, 3, kernel_size=3, stride=2, padding=1, output_padding=1)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5*logvar)
        eps = torch.randn_like(std)
        return mu + eps*std

    def encode(self, x):
        x = self.act(self.conv1(x))
        x = self.act(self.conv2(x))
        x = self.act(self.conv3(x))
        x = self.act(self.conv4(x))
        x = self.act(self.conv5(x))

        x = x.view(-1, 128 * 2 * 2)
        x = self.act(self.fc1(x))
        return self.fc2_mu(x), self.fc2_logvar(x)

    def decode(self, z):
        z = self.act(self.fc3(z))
        z = z.view(-1, 128, 4, 4)
        z = self.act(self.deconv1(z))
        z = self.act(self.deconv2(z))
        z = self.act(self.deconv3(z))
        z = self.act(self.deconv4(z))
        z = self.deconv5(z)
        return torch.sigmoid(z)

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar, z

class EnsembleClassifier(nn.Module):
    def __init__(self, classCount, zdim = 60):
        super(EnsembleClassifier, self).__init__()

        self.classifier_ = nn.Sequential(
            nn.Dropout(),
            nn.Linear(zdim, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, classCount),
        )

    def classifier(self, z):
        return torch.sigmoid(self.classifier_(z))