def __init__(self, args):
super().__init__()
self.input_size = args.input_size
self.n_z = args.z_size
self.MI_obj = args.MI_obj
if args.dataset_name == 'freyfaces':
h_size = 210
elif args.dataset_name == 'cifar10':
h_size = 384
else:
h_size = 294
self.main_x = nn.Sequential(
Conv2d(self.input_size[0], 32, 7, 1, 3, activation = nn.ELU()),
Conv2d(32, 32, 3, 2, 1, activation = nn.ELU()),
Conv2d(32, 64, 5, 1, 2, activation = nn.ELU()),
Conv2d(64, 64, 3, 2, 1, activation = nn.ELU()),
Conv2d(64, 6, 3, 1, 1, activation = nn.ELU())
)
self.main_z = NonLinear(self.n_z, h_size, activation=None)
self.joint = NonLinear(2 * h_size, 1, activation=None)
self.ma_et = None
# weights initialization
for m in self.modules():
if isinstance(m, nn.Linear):
he_init(m)
def __init__(self, args):
super(Decoder, self).__init__()
self.input_size = args.input_size
self.input_type = args.input_type
self.n_z = args.z_size
self.fc = nn.Sequential(GatedDense(self.n_z, 300),
GatedDense(300, np.prod(self.input_size)))
self.main = nn.Sequential(GatedConv2d(self.input_size[0], 64, 3, 1, 1),
GatedConv2d(64, 64, 3, 1, 1),
GatedConv2d(64, 64, 3, 1, 1),
GatedConv2d(64, 64, 3, 1, 1))
if self.input_type == 'binary':
self.p_x_mean = Conv2d(64, 1, 1, 1, 0, activation=nn.Sigmoid())
elif self.input_type == 'gray' or self.input_type == 'continuous':
self.p_x_mean = Conv2d(64,
self.input_size[0],
1,
1,
0,
activation=nn.Sigmoid())
# weights initialization
for m in self.modules():
if isinstance(m, nn.Linear):
he_init(m)
def __init__(self, args):
super(conv_vae, self).__init__(args)
# encoder: q(z | x)
d = 32
act = F.relu
self.q_z_layers = nn.Sequential(
GatedConv2d(self.args.input_size[0], d, 4, 2, 1, activation=act),
GatedConv2d(d, 2 * d, 4, 2, 1, activation=act),
GatedConv2d(2 * d, 4 * d, 4, 2, 1, activation=act),
GatedConv2d(4 * d, 8 * d, 4, 2, 1, activation=act),
GatedConv2d(8 * d, self.args.z1_size, 4, 1, 0, activation=None))
self.q_z_mean = Linear(self.args.z1_size, self.args.z1_size)
self.q_z_logvar = NonLinear(self.args.z1_size,
self.args.z1_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# decoder: p(x | z)
self.p_x_layers = nn.Sequential(
GatedConvTranspose2d(self.args.z1_size,
8 * d,
4,
1,
0,
activation=act),
GatedConvTranspose2d(8 * d, 4 * d, 4, 2, 1, activation=act),
GatedConvTranspose2d(4 * d, 2 * d, 4, 2, 1, activation=act),
GatedConvTranspose2d(2 * d, d, 4, 2, 1, activation=act),
GatedConvTranspose2d(d, d, 4, 2, 1, activation=None))
if self.args.input_type == 'binary':
self.p_x_mean = Conv2d(d, 1, 1, 1, 0, activation=nn.Sigmoid())
elif self.args.input_type == 'gray' or self.args.input_type == 'continuous':
self.p_x_mean = Conv2d(d,
self.args.input_size[0],
1,
1,
0,
activation=nn.Sigmoid())
self.p_x_logvar = Conv2d(d,
self.args.input_size[0],
1,
1,
0,
activation=nn.Hardtanh(min_val=-4.5,
max_val=0.))
# weights initialization
for m in self.modules():
if isinstance(m, nn.Linear):
he_init(m)
# add pseudo-inputs if VampPrior
if self.args.prior == 'vampprior':
self.add_pseudoinputs()
elif self.args.prior == 'GMM':
self.initialize_GMMparams(Kmog=10, mode='random')
def __init__(self, args):
super(VAE, self).__init__(args)
if self.args.dataset_name == 'freyfaces':
h_size = 210
elif self.args.dataset_name == 'cifar10':
h_size = 384
elif self.args.dataset_name == 'celeba':
h_size = 6144
else:
h_size = 294
# encoder: q(z2 | x)
self.q_z2_layers = nn.Sequential(
GatedConv2d(self.args.input_size[0], 32, 7, 1, 3),
GatedConv2d(32, 32, 3, 2, 1), GatedConv2d(32, 64, 5, 1, 2),
GatedConv2d(64, 64, 3, 2, 1), GatedConv2d(64, 6, 3, 1, 1))
# linear layers
self.q_z2_mean = NonLinear(h_size, self.args.z2_size, activation=None)
self.q_z2_logvar = NonLinear(h_size,
self.args.z2_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# encoder: q(z1|x,z2)
# PROCESSING x
self.q_z1_layers_x = nn.Sequential(
GatedConv2d(self.args.input_size[0], 32, 3, 1, 1),
GatedConv2d(32, 32, 3, 2, 1), GatedConv2d(32, 64, 3, 1, 1),
GatedConv2d(64, 64, 3, 2, 1), GatedConv2d(64, 6, 3, 1, 1))
# PROCESSING Z2
self.q_z1_layers_z2 = nn.Sequential(
GatedDense(self.args.z2_size, h_size))
# PROCESSING JOINT
self.q_z1_layers_joint = nn.Sequential(GatedDense(2 * h_size, 300))
# linear layers
self.q_z1_mean = NonLinear(300, self.args.z1_size, activation=None)
self.q_z1_logvar = NonLinear(300,
self.args.z1_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# decoder p(z1|z2)
self.p_z1_layers = nn.Sequential(GatedDense(self.args.z2_size, 300),
GatedDense(300, 300))
self.p_z1_mean = NonLinear(300, self.args.z1_size, activation=None)
self.p_z1_logvar = NonLinear(300,
self.args.z1_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# decoder: p(x | z)
self.p_x_layers_z1 = nn.Sequential(GatedDense(self.args.z1_size, 300))
self.p_x_layers_z2 = nn.Sequential(GatedDense(self.args.z2_size, 300))
self.p_x_layers_joint_pre = nn.Sequential(
GatedDense(2 * 300, np.prod(self.args.input_size)))
# decoder: p(x | z)
act = nn.ReLU(True)
# joint
self.p_x_layers_joint = nn.Sequential(
GatedConv2d(self.args.input_size[0], 64, 3, 1, 1),
GatedConv2d(64, 64, 3, 1, 1),
GatedConv2d(64, 64, 3, 1, 1),
GatedConv2d(64, 64, 3, 1, 1),
)
if self.args.input_type == 'binary':
self.p_x_mean = Conv2d(64, 1, 1, 1, 0, activation=nn.Sigmoid())
elif self.args.input_type == 'gray' or self.args.input_type == 'continuous':
self.p_x_mean = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Sigmoid())
self.p_x_logvar = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Hardtanh(min_val=-4.5,
max_val=0.))
# weights initialization
for m in self.modules():
if isinstance(m, nn.Linear):
he_init(m)
# add pseudo-inputs if VampPrior
if self.args.prior == 'vampprior':
self.add_pseudoinputs()
def create_model(self, args):
if args.dataset_name == 'freyfaces':
self.h_size = 210
elif args.dataset_name == 'cifar10' or args.dataset_name == 'svhn':
self.h_size = 384
else:
self.h_size = 294
fc_size = 300
# encoder: q(z2 | x)
self.q_z_layers = nn.Sequential(
GatedConv2d(self.args.input_size[0],
32,
7,
1,
3,
no_attention=args.no_attention),
GatedConv2d(32, 32, 3, 2, 1, no_attention=args.no_attention),
GatedConv2d(32, 64, 5, 1, 2, no_attention=args.no_attention),
GatedConv2d(64, 64, 3, 2, 1, no_attention=args.no_attention),
GatedConv2d(64, 6, 3, 1, 1, no_attention=args.no_attention))
# linear layers
self.q_z_mean = NonLinear(self.h_size,
self.args.z2_size,
activation=None)
# SAME VARAITIONAL VAR TO SEE IF IT HELPS
self.q_z_logvar = NonLinear(self.h_size,
self.args.z2_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# encoder: q(z1|x,z2)
# PROCESSING x
self.q_z1_layers_x = nn.Sequential(
GatedConv2d(self.args.input_size[0],
32,
3,
1,
1,
no_attention=args.no_attention),
GatedConv2d(32, 32, 3, 2, 1, no_attention=args.no_attention),
GatedConv2d(32, 64, 3, 1, 1, no_attention=args.no_attention),
GatedConv2d(64, 64, 3, 2, 1, no_attention=args.no_attention),
GatedConv2d(64, 6, 3, 1, 1, no_attention=args.no_attention))
# PROCESSING Z2
self.q_z1_layers_z2 = nn.Sequential(
GatedDense(self.args.z2_size, self.h_size))
# PROCESSING JOINT
self.q_z1_layers_joint = nn.Sequential(
GatedDense(2 * self.h_size, fc_size))
# linear layers
self.q_z1_mean = NonLinear(fc_size, self.args.z1_size, activation=None)
self.q_z1_logvar = NonLinear(fc_size,
self.args.z1_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# decoder p(z1|z2)
self.p_z1_layers_z2 = nn.Sequential(
GatedDense(self.args.z2_size,
fc_size,
no_attention=args.no_attention),
GatedDense(fc_size, fc_size, no_attention=args.no_attention))
self.p_z1_mean = NonLinear(fc_size, self.args.z1_size, activation=None)
self.p_z1_logvar = NonLinear(fc_size,
self.args.z1_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# decoder: p(x | z)
self.p_x_layers_z1 = nn.Sequential(
GatedDense(self.args.z1_size,
fc_size,
no_attention=args.no_attention))
self.p_x_layers_z2 = nn.Sequential(
GatedDense(self.args.z2_size,
fc_size,
no_attention=args.no_attention))
self.p_x_layers_joint_pre = nn.Sequential(
GatedDense(2 * fc_size,
np.prod(self.args.input_size),
no_attention=args.no_attention))
# decoder: p(x | z)
self.p_x_layers_joint = nn.Sequential(
GatedConv2d(self.args.input_size[0],
64,
3,
1,
1,
no_attention=args.no_attention),
GatedConv2d(64, 64, 3, 1, 1, no_attention=args.no_attention),
GatedConv2d(64, 64, 3, 1, 1, no_attention=args.no_attention),
GatedConv2d(64, 64, 3, 1, 1, no_attention=args.no_attention),
)
if self.args.input_type == 'binary':
self.p_x_mean = Conv2d(64, 1, 1, 1, 0, activation=nn.Sigmoid())
elif self.args.input_type == 'gray' or self.args.input_type == 'continuous':
self.p_x_mean = Conv2d(64, self.args.input_size[0], 1, 1, 0)
self.p_x_logvar = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Hardtanh(min_val=-4.5,
max_val=0.))
elif self.args.input_type == 'pca':
self.p_x_mean = Conv2d(64, 1, 1, 1, 0)
self.p_x_logvar = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Hardtanh(min_val=-4.5,
max_val=0.))
def __init__(self, args):
super(VAE, self).__init__(args)
if self.args.dataset_name == 'freyfaces':
h_size = 210
elif self.args.dataset_name == 'cifar10':
h_size = 384
else:
h_size = 294
# encoder: q(z2 | x)
self.q_z2_layers = nn.ModuleList()
# conv 0
self.q_z2_layers.append(
GatedConv2d(self.args.input_size[0], 32, 5, 1, 3))
# conv 1
self.q_z2_layers.append(GatedConv2d(32, 32, 5, 2, 1))
# conv 2
self.q_z2_layers.append(GatedConv2d(32, 64, 5, 1, 3))
# conv 3
self.q_z2_layers.append(GatedConv2d(64, 64, 5, 2, 1))
# conv 4
self.q_z2_layers.append(GatedConv2d(64, 6, 3, 1, 1))
# linear layers
self.q_z2_mean = NonLinear(h_size, self.args.z2_size, activation=None)
self.q_z2_logvar = NonLinear(h_size,
self.args.z2_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# encoder: q(z1|x,z2)
self.q_z1_layers_x = nn.ModuleList()
self.q_z1_layers_z2 = nn.ModuleList()
self.q_z1_layers_joint = nn.ModuleList()
# PROCESSING x
# conv 0
self.q_z1_layers_x.append(
GatedConv2d(self.args.input_size[0], 32, 5, 1, 3))
# conv 1
self.q_z1_layers_x.append(GatedConv2d(32, 32, 5, 2, 1))
# conv 2
self.q_z1_layers_x.append(GatedConv2d(32, 64, 5, 1, 3))
# conv 3
self.q_z1_layers_x.append(GatedConv2d(64, 64, 5, 2, 1))
# conv 7
self.q_z1_layers_x.append(GatedConv2d(64, 6, 3, 1, 1))
# PROCESSING Z2
self.q_z1_layers_z2.append(GatedDense(self.args.z2_size, h_size))
# PROCESSING JOINT
self.q_z1_layers_joint.append(GatedDense(2 * h_size, 300))
# linear layers
self.q_z1_mean = NonLinear(300, self.args.z1_size, activation=None)
self.q_z1_logvar = NonLinear(300,
self.args.z1_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# decoder p(z1|z2)
self.p_z1_layers = nn.ModuleList()
self.p_z1_layers.append(GatedDense(self.args.z2_size, 300))
self.p_z1_layers.append(GatedDense(300, 300))
self.p_z1_mean = NonLinear(300, self.args.z1_size, activation=None)
self.p_z1_logvar = NonLinear(300,
self.args.z1_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# decoder: p(x | z)
self.p_x_layers_z1 = nn.ModuleList()
self.p_x_layers_z2 = nn.ModuleList()
self.p_x_layers_z1.append(
GatedDense(self.args.z1_size, np.prod(self.args.input_size)))
self.p_x_layers_z2.append(
GatedDense(self.args.z2_size, np.prod(self.args.input_size)))
# PixelCNN
act = nn.ReLU()
self.pixelcnn = nn.Sequential(
MaskedConv2d('A',
self.args.input_size[0] + 2 * self.args.input_size[0],
64,
5,
1,
2,
bias=True), nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 5, 1, 2, bias=True), nn.BatchNorm2d(64),
act, MaskedConv2d('B', 64, 64, 5, 1, 2, bias=True),
nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 5, 1, 2, bias=True), nn.BatchNorm2d(64),
act, MaskedConv2d('B', 64, 64, 5, 1, 2, bias=True),
nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 5, 1, 2,
bias=True), nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 5, 1, 2,
bias=True), nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 5, 1, 2,
bias=True), nn.BatchNorm2d(64), act)
if self.args.input_type == 'binary':
self.p_x_mean = Conv2d(64, 1, 1, 1, 0, activation=nn.Sigmoid())
elif self.args.input_type == 'gray' or self.args.input_type == 'continuous':
self.p_x_mean = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Sigmoid())
self.p_x_logvar = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Hardtanh(min_val=-4.5,
max_val=0.))
# weights initialization
for m in self.modules():
if isinstance(m, nn.Linear):
he_init(m)
# add pseudo-inputs if VampPrior
if self.args.prior == 'vampprior':
self.add_pseudoinputs()
def __init__(self, args):
super(VAE, self).__init__(args)
h_size = 512
nc = self.args.input_size[0] # number of channels
nz = self.args.z1_size # size of latent vector
ngf = 64 # decoder (generator) filter factor
ndf = 64 # encoder filter factor
self.q_z_layers = nn.Sequential(
# input is (nc) x 64 x 64
nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
nn.LeakyReLU(0.2, inplace=True),
# state size. (ndf) x 32 x 32
nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
nn.BatchNorm2d(ndf * 2),
nn.LeakyReLU(0.2, inplace=True),
# state size. (ndf*2) x 16 x 16
nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
nn.BatchNorm2d(ndf * 4),
nn.LeakyReLU(0.2, inplace=True),
# state size. (ndf*4) x 8 x 8
nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
nn.BatchNorm2d(ndf * 8),
nn.LeakyReLU(0.2, inplace=True),
# state size. (ndf*8) x 4 x 4
nn.Conv2d(ndf * 8, ndf * 8, 4, 1, 0, bias=False))
# linear layers
self.q_z_mean = NonLinear(h_size, self.args.z1_size, activation=None)
self.q_z_logvar = NonLinear(h_size,
self.args.z1_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# decoder: p(x | z)
self.p_x_layers_z = nn.Sequential(GatedDense(self.args.z1_size,
h_size))
# decoder: p(x | z)
act = nn.ReLU(True)
# joint
self.p_x_layers_joint = nn.Sequential(
# input is Z, going into a convolution
nn.ConvTranspose2d(h_size, ngf * 8, 4, 1, 0, bias=False),
nn.BatchNorm2d(ngf * 8),
nn.ReLU(True),
# state size. (ngf*8) x 4 x 4
nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
nn.BatchNorm2d(ngf * 4),
nn.ReLU(True),
# state size. (ngf*4) x 8 x 8
nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
nn.BatchNorm2d(ngf * 2),
nn.ReLU(True),
# state size. (ngf*2) x 16 x 16
nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
nn.BatchNorm2d(ngf),
nn.ReLU(True),
# state size. (ngf) x 32 x 32
nn.ConvTranspose2d(ngf, ngf, 4, 2, 1, bias=False),
nn.BatchNorm2d(ngf),
nn.ReLU(True)
# state size. (nc) x 64 x 64
)
if self.args.input_type == 'binary':
self.p_x_mean = Conv2d(64, 1, 1, 1, 0, activation=nn.Sigmoid())
elif self.args.input_type == 'gray' or self.args.input_type == 'continuous':
self.p_x_mean = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Sigmoid())
self.p_x_logvar = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Hardtanh(min_val=-4.5,
max_val=0.))
# weights initialization
for m in self.modules():
if isinstance(m, nn.Linear):
he_init(m)
# add pseudo-inputs if VampPrior
if self.args.prior == 'vampprior':
self.add_pseudoinputs()
def create_model(self, args):
if args.dataset_name == 'freyfaces':
self.h_size = 210
elif args.dataset_name == 'cifar10' or args.dataset_name == 'svhn':
self.h_size = 384
else:
self.h_size = 294
# encoder: q(z2 | x)
self.q_z_layers = nn.Sequential(
GatedConv2d(self.args.input_size[0], 32, 7, 1, 3),
GatedConv2d(32, 32, 3, 2, 1), GatedConv2d(32, 64, 5, 1, 2),
GatedConv2d(64, 64, 3, 2, 1), GatedConv2d(64, 6, 3, 1, 1))
# linear layers
self.q_z_mean = NonLinear(self.h_size,
self.args.z2_size,
activation=None)
self.q_z_logvar = NonLinear(self.h_size,
self.args.z2_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# encoder: q(z1|x,z2)
# PROCESSING x
self.q_z1_layers_x = nn.Sequential(
GatedConv2d(self.args.input_size[0], 32, 3, 1, 1),
GatedConv2d(32, 32, 3, 2, 1), GatedConv2d(32, 64, 3, 1, 1),
GatedConv2d(64, 64, 3, 2, 1), GatedConv2d(64, 6, 3, 1, 1))
# PROCESSING Z2
self.q_z1_layers_z2 = nn.Sequential(
GatedDense(self.args.z2_size, self.h_size))
# PROCESSING JOINT
self.q_z1_layers_joint = nn.Sequential(GatedDense(
2 * self.h_size, 300))
# linear layers
self.q_z1_mean = NonLinear(300, self.args.z1_size, activation=None)
self.q_z1_logvar = NonLinear(300,
self.args.z1_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# decoder p(z1|z2)
self.p_z1_layers_z2 = nn.Sequential(GatedDense(self.args.z2_size, 300),
GatedDense(300, 300))
self.p_z1_mean = NonLinear(300, self.args.z1_size, activation=None)
self.p_z1_logvar = NonLinear(300,
self.args.z1_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# decoder: p(x | z)
self.p_x_layers_z1 = nn.Sequential(
GatedDense(self.args.z1_size, np.prod(self.args.input_size)))
self.p_x_layers_z2 = nn.Sequential(
GatedDense(self.args.z2_size, np.prod(self.args.input_size)))
# decoder: p(x | z)
act = nn.ReLU(True)
#self.pixelcnn = nn.Sequential(
# MaskedConv2d('A', self.args.input_size[0] + 2 * self.args.input_size[0], 64, 3, 1, 1, bias=False),
# nn.BatchNorm2d(64), act,
# MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False), nn.BatchNorm2d(64), act,
# MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False), nn.BatchNorm2d(64), act,
# MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False), nn.BatchNorm2d(64), act,
# MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False), nn.BatchNorm2d(64), act,
# MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False), nn.BatchNorm2d(64), act,
# MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False), nn.BatchNorm2d(64), act,
# MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False), nn.BatchNorm2d(64), act
#)
self.pixelcnn = PixelSNAIL([28, 28], 64, 64, 3, 1, 4, 64)
if self.args.input_type == 'binary':
self.p_x_mean = Conv2d(64, 1, 1, 1, 0, activation=nn.Sigmoid())
elif self.args.input_type == 'gray' or self.args.input_type == 'continuous':
self.p_x_mean = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Sigmoid(),
bias=False)
self.p_x_logvar = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Hardtanh(min_val=-4.5,
max_val=0.),
bias=False)
def __init__(self, args):
super(VAE, self).__init__(args)
# encoder: q(z | x)
if self.args.dataset_name == 'freyfaces':
h_size = 210
elif self.args.dataset_name == 'cifar10':
h_size = 384
elif self.args.dataset_name == 'svhn':
h_size = 384
else:
h_size = 294
# encoder: q(z2 | x)
self.q_z_layers = nn.Sequential(
Conv2d(self.args.input_size[0], 32, 7, 1, 3, activation=nn.ReLU()),
nn.BatchNorm2d(32), Conv2d(32, 32, 3, 2, 1, activation=nn.ReLU()),
nn.BatchNorm2d(32), Conv2d(32, 64, 5, 1, 2, activation=nn.ReLU()),
nn.BatchNorm2d(64), Conv2d(64, 64, 3, 2, 1, activation=nn.ReLU()),
nn.BatchNorm2d(64), Conv2d(64, 6, 3, 1, 1, activation=nn.ReLU()))
'''
self.q_z_layers = [NonGatedDense(np.prod(self.args.input_size), 300, activation=nn.ReLU())]
for i in range(args.number_hidden):
self.q_z_layers.append(NonGatedDense(300, 300, activation=nn.ReLU()))
self.q_z_layers = nn.ModuleList(self.q_z_layers)
'''
self.q_z_mean = Linear(h_size, self.args.z1_size)
self.q_z_logvar = Linear(
h_size, self.args.z1_size
) #NonLinear(300, self.args.z1_size, activation=nn.Hardtanh(min_val=-6.,max_val=2.))
# decoder: p(x | z)
self.p_x_layers = [
NonGatedDense(self.args.z1_size, 300, activation=nn.ReLU())
]
for i in range(args.number_hidden):
self.p_x_layers.append(
NonGatedDense(300, 300, activation=nn.ReLU()))
self.p_x_layers.append(
NonGatedDense(300,
np.prod(self.args.input_size),
activation=nn.ReLU()))
self.p_x_layers = nn.ModuleList(self.p_x_layers)
# PixelCNN
act = nn.ReLU(True)
self.pixelcnn = nn.Sequential(
MaskedConv2d('A',
self.args.input_size[0] + self.args.input_size[0],
64,
3,
1,
1,
bias=False), nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False), nn.BatchNorm2d(64),
act, MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False),
nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False), nn.BatchNorm2d(64),
act, MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False),
nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 3, 1, 1,
bias=False), nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 3, 1, 1,
bias=False), nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 3, 1, 1,
bias=False), nn.BatchNorm2d(64), act)
if self.args.input_type == 'binary':
self.p_x_mean = Conv2d(64, 1, 1, 1, 0, activation=nn.Sigmoid())
elif self.args.input_type == 'gray' or self.args.input_type == 'continuous':
self.p_x_mean = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Sigmoid(),
bias=False)
self.p_x_logvar = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Hardtanh(min_val=-4.5,
max_val=0.),
bias=False)
# weights initialization
for m in self.modules():
if isinstance(m, nn.Linear):
he_init(m)
else:
if torch.is_tensor(m):
torch.nn.init.kaiming_normal(m)
# add pseudo-inputs if VampPrior
if self.args.prior == 'vampprior':
self.add_pseudoinputs()
def __init__(self, args):
super(VAE, self).__init__(args)
if self.args.dataset_name == 'freyfaces':
h_size = 210
elif self.args.dataset_name.startswith(
'cifar10'
) or self.args.dataset_name == 'coil20' or self.args.dataset_name == 'svhn':
h_size = 384
elif self.args.dataset_name == 'usps':
h_size = 96
elif args.dataset_name == 'celeba':
h_size = 1536
else:
h_size = 294
# encoder: q(z | x)
self.q_z_layers = nn.Sequential(
GatedConv2d(self.args.input_size[0], 32, 7, 1, 3),
GatedConv2d(32, 32, 3, 2, 1), GatedConv2d(32, 64, 5, 1, 2),
GatedConv2d(64, 64, 3, 2, 1), GatedConv2d(64, 6, 3, 1, 1))
self.q_z_mean = NonLinear(h_size, self.args.z1_size, activation=None)
self.q_z_logvar = NonLinear(h_size,
self.args.z1_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
self.fc = nn.Sequential(GatedDense(self.args.z1_size, 300),
GatedDense(300, np.prod(self.args.input_size)))
self.p_x_layers = nn.Sequential(
GatedConv2d(self.args.input_size[0], 64, 3, 1, 1),
GatedConv2d(64, 64, 3, 1, 1),
GatedConv2d(64, 64, 3, 1, 1),
GatedConv2d(64, 64, 3, 1, 1),
)
if self.args.input_type == 'binary':
self.p_x_mean = Conv2d(64, 1, 1, 1, 0, activation=nn.Sigmoid())
elif self.args.input_type == 'gray' or self.args.input_type == 'continuous':
self.p_x_mean = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Sigmoid())
self.p_x_logvar = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Hardtanh(min_val=-4.5,
max_val=0.))
# weights initialization
for m in self.modules():
if isinstance(m, nn.Linear):
he_init(m)
# add pseudo-inputs if VampPrior
if self.args.prior == 'vampprior':
self.add_pseudoinputs()
def __init__(self, args):
super(VAE, self).__init__(args)
# encoder: q(z | x)
self.q_z_layers = iRevNet(nBlocks=[18, 18],
nStrides=[1, 2],
nChannels=[16, 64],
nClasses=300,
init_ds=2,
dropout_rate=0.1,
affineBN=True,
in_shape=self.args.input_size,
mult=4)
#self.q_z_layers = nn.DataParallel(self.q_z_layers)
self.q_z_mean = Linear(300, self.args.z1_size)
self.q_z_logvar = NonLinear(300,
self.args.z1_size,
activation=nn.Hardtanh(min_val=-6.,
max_val=2.))
# decoder: p(x | z)
self.p_x_layers = [
NonGatedDense(self.args.z1_size, 300, activation=nn.ReLU())
]
for i in range(args.number_hidden):
self.p_x_layers.append(
NonGatedDense(300, 300, activation=nn.ReLU()))
self.p_x_layers.append(
NonGatedDense(300,
np.prod(self.args.input_size),
activation=nn.ReLU()))
self.p_x_layers = nn.ModuleList(self.p_x_layers)
# PixelCNN
act = nn.ReLU(True)
self.pixelcnn = nn.Sequential(
MaskedConv2d('A',
self.args.input_size[0] + self.args.input_size[0],
64,
3,
1,
1,
bias=False), nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False), nn.BatchNorm2d(64),
act, MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False),
nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False), nn.BatchNorm2d(64),
act, MaskedConv2d('B', 64, 64, 3, 1, 1, bias=False),
nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 3, 1, 1,
bias=False), nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 3, 1, 1,
bias=False), nn.BatchNorm2d(64), act,
MaskedConv2d('B', 64, 64, 3, 1, 1,
bias=False), nn.BatchNorm2d(64), act)
if self.args.input_type == 'binary':
self.p_x_mean = Conv2d(64, 1, 1, 1, 0, activation=nn.Sigmoid())
elif self.args.input_type == 'gray' or self.args.input_type == 'continuous':
self.p_x_mean = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Sigmoid(),
bias=False)
self.p_x_logvar = Conv2d(64,
self.args.input_size[0],
1,
1,
0,
activation=nn.Hardtanh(min_val=-4.5,
max_val=0.),
bias=False)
# weights initialization
for m in self.modules():
if isinstance(m, nn.Linear):
he_init(m)
# add pseudo-inputs if VampPrior
if self.args.prior == 'vampprior':
self.add_pseudoinputs()
