#base : https://github.com/pianomania/infoGAN-pytorch
import torch.nn as nn


class FrontEnd(nn.Module):
    ''' front end part of discriminator and Q'''

    def __init__(self):
        super(FrontEnd, self).__init__()

        self.main = nn.Sequential(
            nn.Conv2d(1, 64, 4, 2, 1),
            nn.LeakyReLU(0.1, inplace=True),
            nn.Conv2d(64, 128, 4, 2, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.1, inplace=True),
            nn.Conv2d(128, 1024, 7, bias=False),
            nn.BatchNorm2d(1024),
            nn.LeakyReLU(0.1, inplace=True),
        )

    def forward(self, x):
        output = self.main(x)
        return output


class D(nn.Module):

    def __init__(self):
        super(D, self).__init__()

        self.main = nn.Sequential(
            nn.Conv2d(1024, 1, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        output = self.main(x).view(-1, 1)
        return output


class Q(nn.Module):

    def __init__(self, nc):
        super(Q, self).__init__()

        self.conv = nn.Conv2d(1024, 128, 1, bias=False)
        self.bn = nn.BatchNorm2d(128)
        self.lReLU = nn.LeakyReLU(0.1, inplace=True)
        self.conv_mu = nn.Conv2d(128, nc, 1)
        self.conv_var = nn.Conv2d(128, nc, 1)
        self.tanh = nn.Tanh()
    def forward(self, x):
        y = self.conv(x)


        mu = self.tanh(self.conv_mu(y).squeeze())

        var = self.conv_var(y).squeeze().exp()

        return mu, var


class G(nn.Module):

    def __init__(self):
        super(G, self).__init__()

        self.main = nn.Sequential(
            nn.ConvTranspose2d(64, 1024, 1, 1, bias=False),
            nn.BatchNorm2d(1024),
            nn.ReLU(True),
            nn.ConvTranspose2d(1024, 128, 7, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.ConvTranspose2d(64, 1, 4, 2, 1, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        output = self.main(x)
        return output

class E(nn.Module):
    ''' Encoder '''

    def __init__(self, nz, type = 'VAE'):
        super(E, self).__init__()
        self.type = type
        self.main = nn.Sequential(
            nn.Conv2d(1, 64, 4, 2, 1),
            nn.LeakyReLU(0.1, inplace=True),
            nn.Conv2d(64, 128, 4, 2, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.1, inplace=True),
            nn.Conv2d(128, 128, 7, bias=False),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.1, inplace=True),
            nn.Conv2d(128, nz, 1)
        )
        self.fc_mu = nn.Conv2d(nz, nz, 1)
        self.fc_sig = nn.Conv2d(nz, nz, 1)

    def forward(self, x):
       if self.type == 'VAE':
            # VAE
            z_ = self.main(x)
            mu = self.fc_mu(z_)
            logvar = self.fc_sig(z_)
            return mu, logvar
       else:
            # AE
            output = self.main(x).view(-1, 1)
            return output

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)


aa7bbc0717a3b87d4dc604f857592b6989cb21a7
<mask token>


class D(nn.Module):

    def __init__(self):
        super(D, self).__init__()
        self.main = nn.Sequential(nn.Conv2d(1024, 1, 1), nn.Sigmoid())

    def forward(self, x):
        output = self.main(x).view(-1, 1)
        return output


class Q(nn.Module):

    def __init__(self, nc):
        super(Q, self).__init__()
        self.conv = nn.Conv2d(1024, 128, 1, bias=False)
        self.bn = nn.BatchNorm2d(128)
        self.lReLU = nn.LeakyReLU(0.1, inplace=True)
        self.conv_mu = nn.Conv2d(128, nc, 1)
        self.conv_var = nn.Conv2d(128, nc, 1)
        self.tanh = nn.Tanh()

    def forward(self, x):
        y = self.conv(x)
        mu = self.tanh(self.conv_mu(y).squeeze())
        var = self.conv_var(y).squeeze().exp()
        return mu, var


class G(nn.Module):

    def __init__(self):
        super(G, self).__init__()
        self.main = nn.Sequential(nn.ConvTranspose2d(64, 1024, 1, 1, bias=
            False), nn.BatchNorm2d(1024), nn.ReLU(True), nn.ConvTranspose2d
            (1024, 128, 7, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(
            True), nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False), nn.
            BatchNorm2d(64), nn.ReLU(True), nn.ConvTranspose2d(64, 1, 4, 2,
            1, bias=False), nn.Sigmoid())

    def forward(self, x):
        output = self.main(x)
        return output


class E(nn.Module):
    """ Encoder """

    def __init__(self, nz, type='VAE'):
        super(E, self).__init__()
        self.type = type
        self.main = nn.Sequential(nn.Conv2d(1, 64, 4, 2, 1), nn.LeakyReLU(
            0.1, inplace=True), nn.Conv2d(64, 128, 4, 2, 1, bias=False), nn
            .BatchNorm2d(128), nn.LeakyReLU(0.1, inplace=True), nn.Conv2d(
            128, 128, 7, bias=False), nn.BatchNorm2d(128), nn.LeakyReLU(0.1,
            inplace=True), nn.Conv2d(128, nz, 1))
        self.fc_mu = nn.Conv2d(nz, nz, 1)
        self.fc_sig = nn.Conv2d(nz, nz, 1)

    def forward(self, x):
        if self.type == 'VAE':
            z_ = self.main(x)
            mu = self.fc_mu(z_)
            logvar = self.fc_sig(z_)
            return mu, logvar
        else:
            output = self.main(x).view(-1, 1)
            return output


<mask token>

<mask token>


class FrontEnd(nn.Module):
    <mask token>
    <mask token>
    <mask token>


class D(nn.Module):

    def __init__(self):
        super(D, self).__init__()
        self.main = nn.Sequential(nn.Conv2d(1024, 1, 1), nn.Sigmoid())

    def forward(self, x):
        output = self.main(x).view(-1, 1)
        return output


class Q(nn.Module):

    def __init__(self, nc):
        super(Q, self).__init__()
        self.conv = nn.Conv2d(1024, 128, 1, bias=False)
        self.bn = nn.BatchNorm2d(128)
        self.lReLU = nn.LeakyReLU(0.1, inplace=True)
        self.conv_mu = nn.Conv2d(128, nc, 1)
        self.conv_var = nn.Conv2d(128, nc, 1)
        self.tanh = nn.Tanh()

    def forward(self, x):
        y = self.conv(x)
        mu = self.tanh(self.conv_mu(y).squeeze())
        var = self.conv_var(y).squeeze().exp()
        return mu, var


class G(nn.Module):

    def __init__(self):
        super(G, self).__init__()
        self.main = nn.Sequential(nn.ConvTranspose2d(64, 1024, 1, 1, bias=
            False), nn.BatchNorm2d(1024), nn.ReLU(True), nn.ConvTranspose2d
            (1024, 128, 7, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(
            True), nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False), nn.
            BatchNorm2d(64), nn.ReLU(True), nn.ConvTranspose2d(64, 1, 4, 2,
            1, bias=False), nn.Sigmoid())

    def forward(self, x):
        output = self.main(x)
        return output


class E(nn.Module):
    """ Encoder """

    def __init__(self, nz, type='VAE'):
        super(E, self).__init__()
        self.type = type
        self.main = nn.Sequential(nn.Conv2d(1, 64, 4, 2, 1), nn.LeakyReLU(
            0.1, inplace=True), nn.Conv2d(64, 128, 4, 2, 1, bias=False), nn
            .BatchNorm2d(128), nn.LeakyReLU(0.1, inplace=True), nn.Conv2d(
            128, 128, 7, bias=False), nn.BatchNorm2d(128), nn.LeakyReLU(0.1,
            inplace=True), nn.Conv2d(128, nz, 1))
        self.fc_mu = nn.Conv2d(nz, nz, 1)
        self.fc_sig = nn.Conv2d(nz, nz, 1)

    def forward(self, x):
        if self.type == 'VAE':
            z_ = self.main(x)
            mu = self.fc_mu(z_)
            logvar = self.fc_sig(z_)
            return mu, logvar
        else:
            output = self.main(x).view(-1, 1)
            return output


<mask token>

<mask token>


class FrontEnd(nn.Module):
    <mask token>

    def __init__(self):
        super(FrontEnd, self).__init__()
        self.main = nn.Sequential(nn.Conv2d(1, 64, 4, 2, 1), nn.LeakyReLU(
            0.1, inplace=True), nn.Conv2d(64, 128, 4, 2, 1, bias=False), nn
            .BatchNorm2d(128), nn.LeakyReLU(0.1, inplace=True), nn.Conv2d(
            128, 1024, 7, bias=False), nn.BatchNorm2d(1024), nn.LeakyReLU(
            0.1, inplace=True))
    <mask token>


class D(nn.Module):

    def __init__(self):
        super(D, self).__init__()
        self.main = nn.Sequential(nn.Conv2d(1024, 1, 1), nn.Sigmoid())

    def forward(self, x):
        output = self.main(x).view(-1, 1)
        return output


class Q(nn.Module):

    def __init__(self, nc):
        super(Q, self).__init__()
        self.conv = nn.Conv2d(1024, 128, 1, bias=False)
        self.bn = nn.BatchNorm2d(128)
        self.lReLU = nn.LeakyReLU(0.1, inplace=True)
        self.conv_mu = nn.Conv2d(128, nc, 1)
        self.conv_var = nn.Conv2d(128, nc, 1)
        self.tanh = nn.Tanh()

    def forward(self, x):
        y = self.conv(x)
        mu = self.tanh(self.conv_mu(y).squeeze())
        var = self.conv_var(y).squeeze().exp()
        return mu, var


class G(nn.Module):

    def __init__(self):
        super(G, self).__init__()
        self.main = nn.Sequential(nn.ConvTranspose2d(64, 1024, 1, 1, bias=
            False), nn.BatchNorm2d(1024), nn.ReLU(True), nn.ConvTranspose2d
            (1024, 128, 7, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(
            True), nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False), nn.
            BatchNorm2d(64), nn.ReLU(True), nn.ConvTranspose2d(64, 1, 4, 2,
            1, bias=False), nn.Sigmoid())

    def forward(self, x):
        output = self.main(x)
        return output


class E(nn.Module):
    """ Encoder """

    def __init__(self, nz, type='VAE'):
        super(E, self).__init__()
        self.type = type
        self.main = nn.Sequential(nn.Conv2d(1, 64, 4, 2, 1), nn.LeakyReLU(
            0.1, inplace=True), nn.Conv2d(64, 128, 4, 2, 1, bias=False), nn
            .BatchNorm2d(128), nn.LeakyReLU(0.1, inplace=True), nn.Conv2d(
            128, 128, 7, bias=False), nn.BatchNorm2d(128), nn.LeakyReLU(0.1,
            inplace=True), nn.Conv2d(128, nz, 1))
        self.fc_mu = nn.Conv2d(nz, nz, 1)
        self.fc_sig = nn.Conv2d(nz, nz, 1)

    def forward(self, x):
        if self.type == 'VAE':
            z_ = self.main(x)
            mu = self.fc_mu(z_)
            logvar = self.fc_sig(z_)
            return mu, logvar
        else:
            output = self.main(x).view(-1, 1)
            return output


<mask token>

<mask token>


class FrontEnd(nn.Module):
    """ front end part of discriminator and Q"""

    def __init__(self):
        super(FrontEnd, self).__init__()
        self.main = nn.Sequential(nn.Conv2d(1, 64, 4, 2, 1), nn.LeakyReLU(
            0.1, inplace=True), nn.Conv2d(64, 128, 4, 2, 1, bias=False), nn
            .BatchNorm2d(128), nn.LeakyReLU(0.1, inplace=True), nn.Conv2d(
            128, 1024, 7, bias=False), nn.BatchNorm2d(1024), nn.LeakyReLU(
            0.1, inplace=True))

    def forward(self, x):
        output = self.main(x)
        return output


class D(nn.Module):

    def __init__(self):
        super(D, self).__init__()
        self.main = nn.Sequential(nn.Conv2d(1024, 1, 1), nn.Sigmoid())

    def forward(self, x):
        output = self.main(x).view(-1, 1)
        return output


class Q(nn.Module):

    def __init__(self, nc):
        super(Q, self).__init__()
        self.conv = nn.Conv2d(1024, 128, 1, bias=False)
        self.bn = nn.BatchNorm2d(128)
        self.lReLU = nn.LeakyReLU(0.1, inplace=True)
        self.conv_mu = nn.Conv2d(128, nc, 1)
        self.conv_var = nn.Conv2d(128, nc, 1)
        self.tanh = nn.Tanh()

    def forward(self, x):
        y = self.conv(x)
        mu = self.tanh(self.conv_mu(y).squeeze())
        var = self.conv_var(y).squeeze().exp()
        return mu, var


class G(nn.Module):

    def __init__(self):
        super(G, self).__init__()
        self.main = nn.Sequential(nn.ConvTranspose2d(64, 1024, 1, 1, bias=
            False), nn.BatchNorm2d(1024), nn.ReLU(True), nn.ConvTranspose2d
            (1024, 128, 7, 1, bias=False), nn.BatchNorm2d(128), nn.ReLU(
            True), nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False), nn.
            BatchNorm2d(64), nn.ReLU(True), nn.ConvTranspose2d(64, 1, 4, 2,
            1, bias=False), nn.Sigmoid())

    def forward(self, x):
        output = self.main(x)
        return output


class E(nn.Module):
    """ Encoder """

    def __init__(self, nz, type='VAE'):
        super(E, self).__init__()
        self.type = type
        self.main = nn.Sequential(nn.Conv2d(1, 64, 4, 2, 1), nn.LeakyReLU(
            0.1, inplace=True), nn.Conv2d(64, 128, 4, 2, 1, bias=False), nn
            .BatchNorm2d(128), nn.LeakyReLU(0.1, inplace=True), nn.Conv2d(
            128, 128, 7, bias=False), nn.BatchNorm2d(128), nn.LeakyReLU(0.1,
            inplace=True), nn.Conv2d(128, nz, 1))
        self.fc_mu = nn.Conv2d(nz, nz, 1)
        self.fc_sig = nn.Conv2d(nz, nz, 1)

    def forward(self, x):
        if self.type == 'VAE':
            z_ = self.main(x)
            mu = self.fc_mu(z_)
            logvar = self.fc_sig(z_)
            return mu, logvar
        else:
            output = self.main(x).view(-1, 1)
            return output


def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

#base : https://github.com/pianomania/infoGAN-pytorch
import torch.nn as nn


class FrontEnd(nn.Module):
    ''' front end part of discriminator and Q'''

    def __init__(self):
        super(FrontEnd, self).__init__()

        self.main = nn.Sequential(
            nn.Conv2d(1, 64, 4, 2, 1),
            nn.LeakyReLU(0.1, inplace=True),
            nn.Conv2d(64, 128, 4, 2, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.1, inplace=True),
            nn.Conv2d(128, 1024, 7, bias=False),
            nn.BatchNorm2d(1024),
            nn.LeakyReLU(0.1, inplace=True),
        )

    def forward(self, x):
        output = self.main(x)
        return output


class D(nn.Module):

    def __init__(self):
        super(D, self).__init__()

        self.main = nn.Sequential(
            nn.Conv2d(1024, 1, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        output = self.main(x).view(-1, 1)
        return output


class Q(nn.Module):

    def __init__(self, nc):
        super(Q, self).__init__()

        self.conv = nn.Conv2d(1024, 128, 1, bias=False)
        self.bn = nn.BatchNorm2d(128)
        self.lReLU = nn.LeakyReLU(0.1, inplace=True)
        self.conv_mu = nn.Conv2d(128, nc, 1)
        self.conv_var = nn.Conv2d(128, nc, 1)
        self.tanh = nn.Tanh()
    def forward(self, x):
        y = self.conv(x)


        mu = self.tanh(self.conv_mu(y).squeeze())

        var = self.conv_var(y).squeeze().exp()

        return mu, var


class G(nn.Module):

    def __init__(self):
        super(G, self).__init__()

        self.main = nn.Sequential(
            nn.ConvTranspose2d(64, 1024, 1, 1, bias=False),
            nn.BatchNorm2d(1024),
            nn.ReLU(True),
            nn.ConvTranspose2d(1024, 128, 7, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.ConvTranspose2d(64, 1, 4, 2, 1, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        output = self.main(x)
        return output

class E(nn.Module):
    ''' Encoder '''

    def __init__(self, nz, type = 'VAE'):
        super(E, self).__init__()
        self.type = type
        self.main = nn.Sequential(
            nn.Conv2d(1, 64, 4, 2, 1),
            nn.LeakyReLU(0.1, inplace=True),
            nn.Conv2d(64, 128, 4, 2, 1, bias=False),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.1, inplace=True),
            nn.Conv2d(128, 128, 7, bias=False),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.1, inplace=True),
            nn.Conv2d(128, nz, 1)
        )
        self.fc_mu = nn.Conv2d(nz, nz, 1)
        self.fc_sig = nn.Conv2d(nz, nz, 1)

    def forward(self, x):
       if self.type == 'VAE':
            # VAE
            z_ = self.main(x)
            mu = self.fc_mu(z_)
            logvar = self.fc_sig(z_)
            return mu, logvar
       else:
            # AE
            output = self.main(x).view(-1, 1)
            return output

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

