20 | 21 | 22 | 23 | ## Example commands 24 | 25 | For original VAE: 26 | 27 | ```pipenv run python main.py ``` 28 | 29 | To also make figures (reconstruction, samples): 30 | 31 | ```pipenv run python main.py --figs ``` 32 | 33 | For IWAE with 5 importance samples: 34 | 35 | ```pipenv run python main.py --importance_num=5 ``` 36 | 37 | For MIWAE(16, 4): 38 | 39 | ```pipenv run python main.py --mean_num=16 --importance_num=4 ``` 40 | 41 | See [the config file](https://github.com/yoonholee/pytorch-generative/blob/master/utils/config.py) for more options. 42 |
43 | 44 | 45 | 46 | ## Results 47 |
48 | 
49 |
50 | ### Quantitative results on dynamically binarized MNIST
51 | | Method | NLL (this repo) | NLL ([IWAE paper](https://arxiv.org/abs/1509.00519)) | NLL ([MIWAE paper](https://arxiv.org/abs/1802.04537)) | comments |
52 | | ------------- | ------------- | ------------- | ------------- | ---- |
53 | | VAE | 87.01 | 86.76 | - |
54 | | MIWAE(5, 1) | 86.45 | 86.47 | - | listed as VAE with k=5
55 | | MIWAE(1, 5) | 85.18 | 85.54 | - | listed as IWAE with k=5
56 | | MIWAE(64, 1) | 86.07 | - | 86.21 | listed as VAE
57 | | MIWAE(16, 4) | 84.99 | - | - |
58 | | MIWAE(8, 8) | 84.69 | - | 84.97 |
59 | | MIWAE(4, 16) | 84.52 | - | 84.56 |
60 | | MIWAE(1, 64) | 84.37 | - | 84.52 | listed as IWAE
61 |
--------------------------------------------------------------------------------
/assets/mnist_recon.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/yoonholee/pytorch-vae/9dc44aae64f0e2896427ce955a48733d6315bb2d/assets/mnist_recon.gif
--------------------------------------------------------------------------------
/assets/mnist_samples.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/yoonholee/pytorch-vae/9dc44aae64f0e2896427ce955a48733d6315bb2d/assets/mnist_samples.gif
--------------------------------------------------------------------------------
/assets/omni_recon.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/yoonholee/pytorch-vae/9dc44aae64f0e2896427ce955a48733d6315bb2d/assets/omni_recon.gif
--------------------------------------------------------------------------------
/assets/omni_samples.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/yoonholee/pytorch-vae/9dc44aae64f0e2896427ce955a48733d6315bb2d/assets/omni_samples.gif
--------------------------------------------------------------------------------
/data_loader/cifar10.py:
--------------------------------------------------------------------------------
1 | from torchvision import datasets
2 |
3 | class cifar10(datasets.CIFAR10):
4 | def get_mean_img(self):
5 | return self.train_data.mean(0)
6 |
7 |
--------------------------------------------------------------------------------
/data_loader/data_loader.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torchvision import transforms
3 | from .stoch_mnist import stochMNIST
4 | from .omniglot import omniglot
5 | from .fixed_mnist import fixedMNIST
6 | from .cifar10 import cifar10
7 |
8 |
9 | def data_loaders(args):
10 | if args.dataset == 'omniglot':
11 | loader_fn, root = omniglot, './dataset/omniglot'
12 | elif args.dataset == 'fixedmnist':
13 | loader_fn, root = fixedMNIST, './dataset/fixedmnist'
14 | elif args.dataset == 'stochmnist':
15 | loader_fn, root = stochMNIST, './dataset/stochmnist'
16 | elif args.dataset == 'cifar10':
17 | loader_fn, root = cifar10, './dataset/cifar10'
18 |
19 | if args.dataset_dir != '': root = args.dataset_dir
20 | kwargs = {'num_workers': 4, 'pin_memory': True} if args.cuda else {}
21 | train_loader = torch.utils.data.DataLoader(
22 | loader_fn(root, train=True, download=True, transform=transforms.ToTensor()),
23 | batch_size=args.batch_size, shuffle=True, **kwargs)
24 | test_loader = torch.utils.data.DataLoader( # need test bs <=64 to make L_5000 tractable in one pass
25 | loader_fn(root, train=False, download=True, transform=transforms.ToTensor()),
26 | batch_size=args.test_batch_size, shuffle=False, **kwargs)
27 | return train_loader, test_loader
28 |
--------------------------------------------------------------------------------
/data_loader/fixed_mnist.py:
--------------------------------------------------------------------------------
1 | import h5py
2 | import torch
3 | import torch.utils.data as data
4 | from torchvision import transforms
5 | import os
6 | import numpy as np
7 | from PIL import Image
8 | import urllib.request
9 |
10 |
11 | class fixedMNIST(data.Dataset):
12 | """ Binarized MNIST dataset, proposed in
13 | http://proceedings.mlr.press/v15/larochelle11a/larochelle11a.pdf """
14 | train_file = 'binarized_mnist_train.amat'
15 | val_file = 'binarized_mnist_valid.amat'
16 | test_file = 'binarized_mnist_test.amat'
17 |
18 | def __init__(self, root, train=True, transform=None, download=False):
19 | # we ignore transform.
20 | self.root = os.path.expanduser(root)
21 | self.train = train # training set or test set
22 |
23 | if download: self.download()
24 | if not self._check_exists():
25 | raise RuntimeError('Dataset not found.' + ' You can use download=True to download it')
26 |
27 | self.data = self._get_data(train=train)
28 |
29 | def __getitem__(self, index):
30 | img = self.data[index]
31 | img = Image.fromarray(img)
32 | img = transforms.ToTensor()(img).type(torch.FloatTensor)
33 | return img, torch.tensor(-1) # Meaningless tensor instead of target
34 |
35 | def __len__(self):
36 | return len(self.data)
37 |
38 | def _get_data(self, train=True):
39 | with h5py.File(os.path.join(self.root, 'data.h5'), 'r') as hf:
40 | data = hf.get('train' if train else 'test')
41 | data = np.array(data)
42 | return data
43 |
44 | def get_mean_img(self):
45 | return self.data.mean(0).flatten()
46 |
47 | def download(self):
48 | if self._check_exists():
49 | return
50 | if not os.path.exists(self.root):
51 | os.makedirs(self.root)
52 |
53 | print('Downloading MNIST with fixed binarization...')
54 | for dataset in ['train', 'valid', 'test']:
55 | filename = 'binarized_mnist_{}.amat'.format(dataset)
56 | url = 'http://www.cs.toronto.edu/~larocheh/public/datasets/binarized_mnist/binarized_mnist_{}.amat'.format(dataset)
57 | print('Downloading from {}...'.format(url))
58 | local_filename = os.path.join(self.root, filename)
59 | urllib.request.urlretrieve(url, local_filename)
60 | print('Saved to {}'.format(local_filename))
61 |
62 | def filename_to_np(filename):
63 | with open(filename) as f:
64 | lines = f.readlines()
65 | return np.array([[int(i)for i in line.split()] for line in lines]).astype('int8')
66 |
67 | train_data = np.concatenate([filename_to_np(os.path.join(self.root, self.train_file)),
68 | filename_to_np(os.path.join(self.root, self.val_file))])
69 | test_data = filename_to_np(os.path.join(self.root, self.val_file))
70 | with h5py.File(os.path.join(self.root, 'data.h5'), 'w') as hf:
71 | hf.create_dataset('train', data=train_data.reshape(-1, 28, 28))
72 | hf.create_dataset('test', data=test_data.reshape(-1, 28, 28))
73 | print('Done!')
74 |
75 | def _check_exists(self):
76 | return os.path.exists(os.path.join(self.root, 'data.h5'))
77 |
--------------------------------------------------------------------------------
/data_loader/omniglot.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.utils.data as data
3 | from torchvision import transforms
4 | import os
5 | from PIL import Image
6 | import urllib.request
7 | import scipy.io
8 |
9 |
10 | class omniglot(data.Dataset):
11 | """ omniglot dataset """
12 | url = 'https://github.com/yburda/iwae/raw/master/datasets/OMNIGLOT/chardata.mat'
13 |
14 | def __init__(self, root, train=True, transform=None, download=False):
15 | # we ignore transform.
16 | self.root = os.path.expanduser(root)
17 | self.train = train # training set or test set
18 |
19 | if download: self.download()
20 | if not self._check_exists():
21 | raise RuntimeError('Dataset not found. You can use download=True to download it')
22 |
23 | self.data = self._get_data(train=train)
24 |
25 | def __getitem__(self, index):
26 | img = self.data[index].reshape(28, 28)
27 | img = Image.fromarray(img)
28 | img = transforms.ToTensor()(img).type(torch.FloatTensor)
29 | img = torch.bernoulli(img) # stochastically binarize
30 | return img, torch.tensor(-1) # Meaningless tensor instead of target
31 |
32 | def __len__(self):
33 | return len(self.data)
34 |
35 | def _get_data(self, train=True):
36 | def reshape_data(data):
37 | return data.reshape((-1, 28, 28)).reshape((-1, 28*28), order='fortran')
38 |
39 | omni_raw = scipy.io.loadmat(os.path.join(self.root, 'chardata.mat'))
40 | data_str = 'data' if train else 'testdata'
41 | data = reshape_data(omni_raw[data_str].T.astype('float32'))
42 | return data
43 |
44 | def get_mean_img(self):
45 | return self.data.mean(0)
46 |
47 | def download(self):
48 | if self._check_exists():
49 | return
50 | if not os.path.exists(self.root):
51 | os.makedirs(self.root)
52 |
53 | print('Downloading from {}...'.format(self.url))
54 | local_filename = os.path.join(self.root, 'chardata.mat')
55 | urllib.request.urlretrieve(self.url, local_filename)
56 | print('Saved to {}'.format(local_filename))
57 |
58 | def _check_exists(self):
59 | return os.path.exists(os.path.join(self.root, 'chardata.mat'))
60 |
--------------------------------------------------------------------------------
/data_loader/stoch_mnist.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torchvision import datasets, transforms
3 | from PIL import Image
4 |
5 |
6 | class stochMNIST(datasets.MNIST):
7 | """ Gets a new stochastic binarization of MNIST at each call. """
8 | def __getitem__(self, index):
9 | if self.train:
10 | img, target = self.train_data[index], self.train_labels[index]
11 | else:
12 | img, target = self.test_data[index], self.test_labels[index]
13 |
14 | img = Image.fromarray(img.numpy(), mode='L')
15 | img = transforms.ToTensor()(img)
16 | img = torch.bernoulli(img) # stochastically binarize
17 | return img, target
18 |
19 | def get_mean_img(self):
20 | imgs = self.train_data.type(torch.float) / 255
21 | mean_img = imgs.mean(0).reshape(-1).numpy()
22 | return mean_img
23 |
--------------------------------------------------------------------------------
/exp.sh:
--------------------------------------------------------------------------------
1 | run="pipenv run python main.py"
2 |
3 | $run --gpu=0 &
4 | $run --gpu=1 --importance_num=64 &
5 | $run --gpu=2 --importance_num=8 --mean_num=8 &
6 | $run --gpu=3 --no_iwae_lr &
7 | $run --gpu=4 --z=100 &
8 | wait
9 |
--------------------------------------------------------------------------------
/main.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 |
4 | import numpy as np
5 | import torch
6 | from tensorboardX import SummaryWriter
7 | from torch import optim
8 |
9 | from data_loader.data_loader import data_loaders
10 | from model.bernoulli_vae import BernoulliVAE
11 | from model.conv_vae import ConvVAE
12 | from utils.config import get_args
13 | from utils.draw_figs import draw_figs
14 |
15 | args = get_args()
16 | os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
17 | args.cuda = torch.cuda.is_available()
18 | device = torch.device("cuda:0" if args.cuda else "cpu")
19 | train_loader, test_loader = data_loaders(args)
20 | torch.manual_seed(args.seed)
21 | if args.cuda:
22 | torch.cuda.manual_seed_all(args.seed)
23 | writer = SummaryWriter(args.out_dir)
24 |
25 | model_class = BernoulliVAE if args.arch == "bernoulli" else ConvVAE
26 | mean_img = train_loader.dataset.get_mean_img()
27 | model = model_class(
28 | device=device,
29 | img_shape=args.img_shape,
30 | h_dim=args.h_dim,
31 | z_dim=args.z_dim,
32 | analytic_kl=args.analytic_kl,
33 | mean_img=mean_img,
34 | ).to(device)
35 | optimizer = optim.Adam(model.parameters(), lr=args.learning_rate, eps=1e-4)
36 | if args.no_iwae_lr:
37 | scheduler = optim.lr_scheduler.ReduceLROnPlateau(
38 | optimizer, mode="min", patience=100, factor=10 ** (-1 / 7)
39 | )
40 | else:
41 | milestones = np.cumsum([3 ** i for i in range(8)])
42 | scheduler = optim.lr_scheduler.MultiStepLR(
43 | optimizer, milestones=milestones, gamma=10 ** (-1 / 7)
44 | )
45 |
46 |
47 | def train(epoch):
48 | for batch_idx, (data, _) in enumerate(train_loader):
49 | optimizer.zero_grad()
50 | outs = model(data, mean_n=args.mean_num, imp_n=args.importance_num)
51 | loss_1, loss = -outs["elbo"].cpu().data.numpy().mean(), outs["loss"].mean()
52 | loss.backward()
53 | optimizer.step()
54 | model.train_step += 1
55 | if model.train_step % args.log_interval == 0:
56 | print(
57 | "Train Epoch: {} ({:.0f}%)\tLoss: {:.6f}".format(
58 | epoch, 100.0 * batch_idx / len(train_loader), loss.item()
59 | )
60 | )
61 | writer.add_scalar("train/loss", loss.item(), model.train_step)
62 | writer.add_scalar("train/loss_1", loss_1, model.train_step)
63 |
64 |
65 | def test(epoch):
66 | elbos = [
67 | model(data, mean_n=1, imp_n=args.log_likelihood_k)["elbo"].squeeze(0)
68 | for data, _ in test_loader
69 | ]
70 |
71 | def get_loss_k(k):
72 | losses = [
73 | model.logmeanexp(elbo[:k], 0).cpu().numpy().flatten() for elbo in elbos
74 | ]
75 | return -np.concatenate(losses).mean()
76 |
77 | return map(get_loss_k, [args.importance_num, 1, 64, args.log_likelihood_k])
78 |
79 |
80 | if args.eval:
81 | model.load_state_dict(torch.load(args.best_model_file))
82 | with torch.no_grad():
83 | print(list(test(0)))
84 | if args.figs:
85 | draw_figs(model, args, test_loader, 0)
86 | sys.exit()
87 |
88 | for epoch in range(1, args.epochs + 1):
89 | writer.add_scalar("learning_rate", optimizer.param_groups[0]["lr"], epoch)
90 | train(epoch)
91 | with torch.no_grad():
92 | if args.figs and epoch % 100 == 1:
93 | draw_figs(model, args, test_loader, epoch)
94 | test_loss, test_1, test_64, test_ll = test(epoch)
95 | if test_loss < model.best_loss:
96 | model.best_loss = test_loss
97 | torch.save(model.state_dict(), args.best_model_file)
98 | scheduler_args = {"metrics": test_loss} if args.no_iwae_lr else {}
99 | scheduler.step(**scheduler_args)
100 | writer.add_scalar("test/loss", test_loss, epoch)
101 | writer.add_scalar("test/loss_1", test_1, epoch)
102 | writer.add_scalar("test/loss_64", test_64, epoch)
103 | writer.add_scalar("test/LL", test_ll, epoch)
104 | print("==== Testing. LL: {:.4f} ====\n".format(test_ll))
105 |
106 | if args.to_gsheets:
107 | from utils.to_sheets import upload_to_google_sheets
108 |
109 | row_data = [args.exp_name, str(test_ll), str(test_64), str(test_64 - test_ll)]
110 | upload_to_google_sheets(row_data=row_data)
111 |
--------------------------------------------------------------------------------
/model/bernoulli_vae.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import torch
3 | from torch import nn
4 | from torch.distributions.bernoulli import Bernoulli
5 | from torch.distributions.normal import Normal
6 |
7 | from .vae_base import VAE
8 |
9 |
10 | class BernoulliVAE(VAE):
11 | def __init__(self, device, img_shape, h_dim, z_dim, analytic_kl, mean_img):
12 | super().__init__(device, z_dim, analytic_kl)
13 | x_dim = np.prod(img_shape)
14 | self.img_shape = img_shape
15 | self.proc_data = lambda x: x.to(device).reshape(-1, x_dim)
16 | self.encoder = nn.Sequential(
17 | nn.Linear(x_dim, h_dim), nn.Tanh(), nn.Linear(h_dim, h_dim), nn.Tanh()
18 | )
19 | self.enc_mu = nn.Linear(h_dim, z_dim)
20 | self.enc_sig = nn.Linear(h_dim, z_dim)
21 | self.decoder = nn.Sequential(
22 | nn.Linear(z_dim, h_dim), nn.Tanh(),
23 | nn.Linear(h_dim, h_dim), nn.Tanh(),
24 | nn.Linear(h_dim, x_dim),
25 | ) # using Bern(logit) is equivalent to putting sigmoid here.
26 |
27 | self.apply(self.init)
28 | mean_img = np.clip(mean_img, 1e-8, 1.0 - 1e-7)
29 | mean_img_logit = np.log(mean_img / (1.0 - mean_img))
30 | self.decoder[-1].bias = torch.nn.Parameter(torch.Tensor(mean_img_logit))
31 |
32 | def init(self, module):
33 | if type(module) == nn.Linear:
34 | torch.nn.init.xavier_uniform_(
35 | module.weight, gain=nn.init.calculate_gain("tanh")
36 | )
37 | module.bias.data.fill_(0.01)
38 |
39 | def encode(self, x):
40 | x = self.proc_data(x)
41 | h = self.encoder(x)
42 | mu, _std = self.enc_mu(h), self.enc_sig(h)
43 | return Normal(mu, nn.functional.softplus(_std)) # torch.exp(.5 * _std)
44 |
45 | def decode(self, z):
46 | x = self.decoder(z)
47 | return Bernoulli(logits=x)
48 |
49 | def lpxz(self, true_x, x_dist):
50 | return x_dist.log_prob(true_x).sum(-1)
51 |
52 | def sample(self, num_samples=64):
53 | z = self.prior.sample((num_samples,))
54 | x_dist = self.decode(z)
55 | return x_dist.sample().view(num_samples, *self.img_shape)
56 |
--------------------------------------------------------------------------------
/model/conv_vae.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from torch import nn
3 | from torch.distributions.normal import Normal
4 | from .vae_base import VAE
5 |
6 |
7 | class Flatten(nn.Module):
8 | def forward(self, input):
9 | return input.view(input.size(0), 16 * 8 * 8).contiguous()
10 |
11 |
12 | class UnFlatten(nn.Module):
13 | def forward(self, input):
14 | return input.view(input.size(0), 16, 8, 8).contiguous()
15 |
16 |
17 | class ConvVAE(VAE):
18 | # XXX: This class does not work at the moment
19 | def __init__(self, device, x_dim, h_dim, z_dim, analytic_kl, mean_img):
20 | # FIXME: integrate so that plot etc works.
21 | VAE.__init__(self, device, x_dim, h_dim, z_dim, analytic_kl, mean_img)
22 | self.proc_data = lambda x: x.to(device)
23 | self.encoder = nn.Sequential(
24 | nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1), nn.ReLU(),
25 | nn.Conv2d(32, 32, kernel_size=3, stride=2, padding=1), nn.ReLU(),
26 | nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1), nn.ReLU(),
27 | nn.Conv2d(32, 16, kernel_size=3, stride=2, padding=1), nn.ReLU(),
28 | Flatten())
29 | self.enc_mu = nn.Linear(16 * 8 * 8, z_dim)
30 | self.enc_sig = nn.Linear(16 * 8 * 8, z_dim)
31 | self.decoder = nn.Sequential(
32 | nn.Linear(z_dim, 16 * 8 * 8), nn.ReLU(),
33 | UnFlatten(),
34 | nn.ConvTranspose2d(16, 32, kernel_size=3, stride=2, padding=1, output_padding=1), nn.ReLU(),
35 | nn.ConvTranspose2d(32, 32, kernel_size=3, stride=1, padding=1), nn.ReLU(),
36 | nn.ConvTranspose2d(32, 32, kernel_size=3, stride=2, padding=1, output_padding=1), nn.ReLU(),
37 | nn.ConvTranspose2d(32, 6, kernel_size=3, stride=1, padding=1))
38 |
39 | self.apply(self.init)
40 | self.decoder[-1].bias = torch.nn.Parameter(torch.cat(
41 | [torch.Tensor(mean_img.mean(0).mean(0)) / 256, .01 * torch.ones([3])]))
42 |
43 | def init(self, module):
44 | if type(module) in [nn.Conv2d, nn.ConvTranspose2d]:
45 | torch.nn.init.xavier_uniform_(module.weight, gain=nn.init.calculate_gain('relu'))
46 | module.bias.data.fill_(.01)
47 |
48 | def encode(self, x):
49 | x = self.proc_data(x)
50 | h = self.encoder(x)
51 | mu, _std = self.enc_mu(h), self.enc_sig(h)
52 | return Normal(mu, nn.functional.softplus(_std))
53 |
54 | def decode(self, z):
55 | mean_n, imp_n, bs = z.size(0), z.size(1), z.size(2)
56 | z = z.view([mean_n * imp_n * bs, -1]).contiguous()
57 | x = self.decoder(z)
58 | x = x.view([mean_n, imp_n, bs, 6, 32, 32]).contiguous()
59 | x_mean, x_std = x[:, :, :, :3, :, :].contiguous(), nn.functional.softplus(x[:, :, :, 3:, :, :]).contiguous()
60 | return Normal(x_mean, x_std)
61 |
62 | def lpxz(self, true_x, x_dist):
63 | return x_dist.log_prob(true_x).sum([-1, -2, -3])
64 |
--------------------------------------------------------------------------------
/model/vae_base.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import torch
3 | from torch import nn
4 | from torch.distributions.normal import Normal
5 |
6 |
7 | class VAE(nn.Module):
8 | def __init__(self, device, z_dim, analytic_kl):
9 | super().__init__()
10 | self.train_step = 0
11 | self.best_loss = np.inf
12 | self.analytic_kl = analytic_kl
13 | self.prior = Normal(
14 | torch.zeros([z_dim]).to(device), torch.ones([z_dim]).to(device)
15 | )
16 |
17 | def proc_data(self, x):
18 | pass
19 |
20 | def encode(self, x):
21 | pass
22 |
23 | def decode(self, z):
24 | pass
25 |
26 | def lpxz(self, true_x, x_dist):
27 | pass
28 |
29 | def sample(self, num_samples=64):
30 | pass
31 |
32 | def elbo(self, true_x, z, x_dist, z_dist):
33 | true_x = self.proc_data(true_x)
34 | lpxz = self.lpxz(true_x, x_dist)
35 |
36 | if self.analytic_kl:
37 | # SGVB^B: -KL(q(z|x)||p(z)) + log p(x|z). Use when KL can be done analytically.
38 | assert z.size(0) == 1 and z.size(1) == 1
39 | kl = torch.distributions.kl.kl_divergence(z_dist, self.prior).sum(-1)
40 | else:
41 | # SGVB^A: log p(z) - log q(z|x) + log p(x|z)
42 | lpz = self.prior.log_prob(z).sum(-1)
43 | lqzx = z_dist.log_prob(z).sum(-1)
44 | kl = -lpz + lqzx
45 | return -kl + lpxz
46 |
47 | def logmeanexp(self, inputs, dim=1):
48 | if inputs.size(dim) == 1:
49 | return inputs
50 | else:
51 | input_max = inputs.max(dim, keepdim=True)[0]
52 | return (inputs - input_max).exp().mean(dim).log() + input_max
53 |
54 | def forward(self, true_x, mean_n, imp_n):
55 | z_dist = self.encode(true_x)
56 | # mean_n, imp_n, batch_size, z_dim
57 | z = z_dist.rsample(torch.Size([mean_n, imp_n]))
58 | x_dist = self.decode(z)
59 |
60 | elbo = self.elbo(true_x, z, x_dist, z_dist) # mean_n, imp_n, batch_size
61 | elbo_iwae = self.logmeanexp(elbo, 1).squeeze(1) # mean_n, batch_size
62 | elbo_iwae_m = torch.mean(elbo_iwae, 0) # batch_size
63 | return {"elbo": elbo, "loss": -elbo_iwae_m}
64 |
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | -i https://pypi.org/simple
2 | certifi==2019.6.16
3 | chardet==3.0.4
4 | cycler==0.10.0
5 | gspread==3.1.0
6 | h5py==2.9.0
7 | httplib2==0.13.1
8 | idna==2.8
9 | imageio==2.5.0
10 | kiwisolver==1.1.0
11 | matplotlib==3.1.1
12 | numpy==1.17.0
13 | oauth2client==4.1.3
14 | pathlib==1.0.1
15 | pillow==6.1.0
16 | protobuf==3.9.1
17 | pyasn1-modules==0.2.6
18 | pyasn1==0.4.6
19 | pyparsing==2.4.2
20 | python-dateutil==2.8.0
21 | requests==2.22.0
22 | rsa==4.0
23 | scipy==1.3.1
24 | six==1.12.0
25 | tensorboardx==1.8
26 | torch==1.2.0
27 | torchvision==0.4.0
28 | urllib3==1.25.3
29 |
--------------------------------------------------------------------------------
/utils/config.py:
--------------------------------------------------------------------------------
1 | import os
2 | import argparse
3 |
4 | parser = argparse.ArgumentParser()
5 | parser.add_argument('--gpu', type=int, default=0)
6 | parser.add_argument('--seed', type=int, default=42)
7 | parser.add_argument('--log_interval', type=int, default=500)
8 | parser.add_argument('--eval', action='store_true')
9 | parser.add_argument('--figs', action='store_true')
10 | parser.add_argument('--to_gsheets', action='store_true')
11 | parser.add_argument('--arch', type=str, default='bernoulli', choices=['bernoulli']) # TODO: make conv work
12 |
13 | parser.add_argument('--dataset_dir', type=str, default='')
14 | parser.add_argument('--dataset', type=str, default='stochmnist',
15 | choices=['stochmnist', 'omniglot', 'fixedmnist']) # TODO: make cifar10 work
16 | parser.add_argument('--batch_size', type=int, default=20) # iwae uses 20
17 | parser.add_argument('--test_batch_size', type=int, default=64)
18 | parser.add_argument('--epochs', type=int, default=3280) # iwae uses 3280
19 |
20 | parser.add_argument('--learning_rate', type=float, default=1e-3)
21 | parser.add_argument('--no_iwae_lr', action='store_true')
22 | parser.add_argument('--mean_num', type=int, default=1) # M in "tighter variational bounds...". Use 1 for vanilla vae
23 | parser.add_argument('--importance_num', type=int, default=1) # k of iwae. Use 1 for vanilla vae
24 | parser.add_argument('--analytic_kl', action='store_true')
25 | parser.add_argument('--h_dim', type=int, default=200)
26 | parser.add_argument('--z_dim', type=int, default=50)
27 |
28 |
29 | def get_args():
30 | args = parser.parse_args()
31 |
32 | def cstr(arg, arg_name, default, custom_str=False):
33 | """ Get config str for arg, ignoring if set to default. """
34 | not_default = arg != default
35 | if not custom_str:
36 | custom_str = f'_{arg_name}{arg}'
37 | return custom_str if not_default else ''
38 |
39 | args.exp_name = (f'm{args.mean_num}_k{args.importance_num}'
40 | f'{cstr(args.dataset, "", "stochmnist")}{cstr(args.arch, "", "bernoulli")}'
41 | f'{cstr(args.seed, "seed", 42)}{cstr(args.batch_size, "bs", 20)}'
42 | f'{cstr(args.h_dim, "h", 200)}{cstr(args.z_dim, "z", 50)}'
43 | f'{cstr(args.learning_rate, "lr", 1e-3)}{cstr(args.analytic_kl, None, False, "_analytic")}'
44 | f'{cstr(args.no_iwae_lr, None, False, "_noiwae")}{cstr(args.epochs, "epoch", 3280)}')
45 |
46 | args.figs_dir = os.path.join('figs', args.exp_name)
47 | args.out_dir = os.path.join('result', args.exp_name)
48 | args.best_model_file = os.path.join('result', args.exp_name, 'best_model.pt')
49 | if not os.path.exists(args.out_dir):
50 | os.makedirs(args.out_dir)
51 | if not os.path.exists(args.figs_dir):
52 | os.makedirs(args.figs_dir)
53 |
54 | args.log_likelihood_k = 100 if args.dataset == 'cifar10' else 5000
55 | args.img_shape = (32, 32) if args.dataset == 'cifar10' else (28, 28)
56 | return args
57 |
--------------------------------------------------------------------------------
/utils/draw_figs.py:
--------------------------------------------------------------------------------
1 | import matplotlib
2 | matplotlib.use('Agg')
3 | import matplotlib.pyplot as plt
4 | import imageio
5 | import pathlib
6 | import numpy as np
7 | import torch
8 |
9 |
10 | def draw_gif(name, figs_dir, glob_str):
11 | files = [file for file in pathlib.Path(figs_dir).glob(glob_str)]
12 | images = [imageio.imread(str(file)) for file in sorted(files)]
13 | imageio.mimsave('{}/{}'.format(figs_dir, name), images, duration=.5)
14 |
15 |
16 | def draw_figs(model, args, test_loader, epoch):
17 | samples = model.sample(num_samples=100).data.cpu().numpy()
18 | plt.figure(figsize=(5, 5))
19 | plt.suptitle('Samples, Epoch {}'.format(epoch), fontsize=20)
20 | plt.axis('square')
21 | plt.legend(frameon=True)
22 | for idx, im in enumerate(samples):
23 | plt.subplot(10, 10, idx+1)
24 | plt.imshow(im, cmap='Greys')
25 | plt.axis('off')
26 | plt.savefig('figs/{}/samples_{:04}.jpg'.format(args.exp_name, epoch))
27 | plt.clf()
28 | draw_gif('{}_samples.gif'.format(args.exp_name), args.figs_dir, 'samples*.jpg')
29 |
30 | for batch_idx, (data, _) in enumerate(test_loader):
31 | break
32 | z_dist = model.encode(data)
33 | z = z_dist.rsample()
34 | recon = model.decode(z).probs.view(args.test_batch_size, 28, 28)
35 | data = data.view(args.test_batch_size, 28, 28)
36 | plt.figure(figsize=(5, 5))
37 | plt.suptitle('Reconstruction, Epoch {}'.format(epoch), fontsize=20)
38 | plt.axis('square')
39 | plt.legend(frameon=True)
40 | for i in range(50):
41 | data_i = data[i].data.cpu().numpy()
42 | recon_i = recon[i].data.cpu().numpy()
43 | plt.subplot(10, 10, 2*i+1)
44 | plt.imshow(data_i, cmap='Greys')
45 | plt.axis('off')
46 | plt.subplot(10, 10, 2*i+2)
47 | plt.imshow(recon_i, cmap='Greys')
48 | plt.axis('off')
49 | plt.savefig('figs/{}/reconstruction_{:04}.jpg'.format(args.exp_name, epoch))
50 | plt.clf()
51 | draw_gif('{}_reconstruction.gif'.format(args.exp_name), args.figs_dir, 'reconstruction*.jpg')
52 |
53 | if args.z_dim == 2:
54 | latent_space, labels = [], []
55 | for batch_idx, (data, label) in enumerate(test_loader):
56 | latent_space.append(model.encode(data).loc.data.cpu().numpy())
57 | labels.append(label)
58 | latent_space, labels = np.concatenate(latent_space), np.concatenate(labels)
59 | plt.figure(figsize=(5, 5))
60 | for c in range(10):
61 | idx = (labels == c)
62 | plt.scatter(latent_space[idx, 0], latent_space[idx, 1],
63 | c=matplotlib.cm.get_cmap('tab10')(c), marker=',', label=str(c), alpha=.7)
64 | plt.suptitle('Latent representation, Epoch {}'.format(epoch), fontsize=20)
65 | plt.axis('square')
66 | plt.legend(frameon=True)
67 | plt.savefig('figs/{}/latent_{:04}.jpg'.format(args.exp_name, epoch))
68 | plt.clf()
69 | draw_gif('{}_latent.gif'.format(args.exp_name), args.figs_dir, 'latent*.jpg')
70 |
71 | plt.close('all')
72 |
--------------------------------------------------------------------------------
/utils/to_sheets.py:
--------------------------------------------------------------------------------
1 | import gspread
2 | from oauth2client.service_account import ServiceAccountCredentials
3 | get_credentials = ServiceAccountCredentials.from_json_keyfile_name
4 |
5 | scope = ['https://spreadsheets.google.com/feeds',
6 | 'https://www.googleapis.com/auth/drive']
7 | # To make this work, obtain credentials from Google Sheets API and save to
8 | # creds.json in current directory.
9 | credentials = get_credentials('creds.json', scope)
10 | gc = gspread.authorize(credentials)
11 | sheet_name = 'pytorch-generative'
12 |
13 |
14 | def upload_to_google_sheets(row_data, index=2):
15 | worksheet = gc.open(sheet_name).sheet1
16 | worksheet.insert_row(row_data, index=index)
17 |
--------------------------------------------------------------------------------