Skip to content

Commit

Permalink
Updated readme
Browse files Browse the repository at this point in the history
  • Loading branch information
@AntixK
AntixK committed Feb 7, 2020
1 parent 3cba112 commit 76b4a6f
Show file tree
Hide file tree
Showing 5 changed files with 54 additions and 108 deletions.
51 changes: 35 additions & 16 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ exp_params:
trainer_params:
gpus: 1
max_nb_epochs: 50
gradient_clip_val: 0.005
gradient_clip_val: 1.5
.
.
.
Expand All @@ -78,21 +78,24 @@ logging_params:
</h2>
| Model | Paper |Reconstruction | Samples |
|-----------------------|--------------------------------------------------|---------------|---------|
| VAE |[Link](https://arxiv.org/abs/1312.6114) | ![][2] | ![][1] |
| Conditional VAE |[Link](https://openreview.net/forum?id=rJWXGDWd-H)| ![][16] | ![][15] |
| WAE - MMD (RBF Kernel)|[Link](https://arxiv.org/abs/1711.01558) | ![][4] | ![][3] |
| WAE - MMD (IMQ Kernel)|[Link](https://arxiv.org/abs/1711.01558) | ![][6] | ![][5] |
| Beta-VAE |[Link](https://openreview.net/forum?id=Sy2fzU9gl) | ![][8] | ![][7] |
| Disentangled Beta-VAE |[Link](https://arxiv.org/abs/1804.03599) | ![][22] | ![][21] |
| IWAE (5 Samples) |[Link](https://arxiv.org/abs/1509.00519) | ![][10] | ![][9] |
| DFCVAE |[Link](https://arxiv.org/abs/1610.00291) | ![][12] | ![][11] |
| MSSIM VAE |[Link](https://arxiv.org/abs/1511.06409) | ![][14] | ![][13] |
| Categorical VAE |[Link](https://arxiv.org/abs/1611.01144) | ![][18] | ![][17] |
| Joint VAE |[Link](https://arxiv.org/abs/1804.00104) | ![][20] | ![][19] |
| Info VAE |[Link](https://arxiv.org/abs/1706.02262) | ![][24] | ![][23] |
| LogCosh VAE |[Link](https://openreview.net/forum?id=rkglvsC9Ym)| ![][26] | ![][25] |
| Model | Paper |Reconstruction | Samples |
|-------------------------------------|--------------------------------------------------|---------------|---------|
| VAE ![][vae_code] |[Link](https://arxiv.org/abs/1312.6114) | ![][2] | ![][1] |
| Conditional VAE ![][cvae_code] |[Link](https://openreview.net/forum?id=rJWXGDWd-H)| ![][16] | ![][15] |
| WAE - MMD (RBF Kernel) ![][wae_code]|[Link](https://arxiv.org/abs/1711.01558) | ![][4] | ![][3] |
| WAE - MMD (IMQ Kernel) ![][wae_code]|[Link](https://arxiv.org/abs/1711.01558) | ![][6] | ![][5] |
| Beta-VAE ![][bvae_code] |[Link](https://openreview.net/forum?id=Sy2fzU9gl) | ![][8] | ![][7] |
| Disentangled Beta-VAE ![][bvae_code]|[Link](https://arxiv.org/abs/1804.03599) | ![][22] | ![][21] |
| IWAE (5 Samples) ![][iwae_code] |[Link](https://arxiv.org/abs/1509.00519) | ![][10] | ![][9] |
| DFCVAE ![][dfcvae_code] |[Link](https://arxiv.org/abs/1610.00291) | ![][12] | ![][11] |
| MSSIM VAE ![][mssimvae_code] |[Link](https://arxiv.org/abs/1511.06409) | ![][14] | ![][13] |
| Categorical VAE ![][catvae_code] |[Link](https://arxiv.org/abs/1611.01144) | ![][18] | ![][17] |
| Joint VAE ![][jointvae_code] |[Link](https://arxiv.org/abs/1804.00104) | ![][20] | ![][19] |
| Info VAE ![][infovae_code] |[Link](https://arxiv.org/abs/1706.02262) | ![][24] | ![][23] |
| LogCosh VAE ![][logcoshvae_code] |[Link](https://openreview.net/forum?id=rkglvsC9Ym)| ![][26] | ![][25] |
| SWAE (50 Projections) ![][swae_code]|[Link](https://arxiv.org/abs/1804.01947) | ![][28] | ![][27] |
<!-- | Gamma VAE |[Link](https://arxiv.org/abs/1610.05683) | ![][16] | ![][15] |-->
Expand Down Expand Up @@ -151,6 +154,20 @@ I would be happy to include your result (along with your config file) in this re
```
-----------

[vae_code]: https://github.com/AntixK/PyTorch-VAE/blob/master/models/vanilla_vae.py
[cvae_code]: https://github.com/AntixK/PyTorch-VAE/blob/master/models/cvae.py
[bvae_code]: https://github.com/AntixK/PyTorch-VAE/blob/master/models/beta_vae.py
[wae_code]: https://github.com/AntixK/PyTorch-VAE/blob/master/models/wae_mmd.py
[iwae_code]: https://github.com/AntixK/PyTorch-VAE/blob/master/models/iwae.py
[swae_code]: https://github.com/AntixK/PyTorch-VAE/blob/master/models/swae.py
[jointvae_code]: https://github.com/AntixK/PyTorch-VAE/blob/master/models/joint_vae.py
[dfcvae_code]: https://github.com/AntixK/PyTorch-VAE/blob/master/models/dfcvae.py
[mssimvae_code]: https://github.com/AntixK/PyTorch-VAE/blob/master/models/mssim_vae.py
[logcoshvae_code]: https://github.com/AntixK/PyTorch-VAE/blob/master/models/logcosh_vae.py
[catvae_code]: https://github.com/AntixK/PyTorch-VAE/blob/master/models/cat_vae.py



[1]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/Vanilla%20VAE_25.png
[2]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/recons_Vanilla%20VAE_25.png
[3]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/WAE_RBF_18.png
Expand All @@ -177,6 +194,8 @@ I would be happy to include your result (along with your config file) in this re
[24]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/recons_InfoVAE_31.png
[25]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/LogCoshVAE_49.png
[26]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/recons_LogCoshVAE_49.png
[27]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/SWAE_49.png
[28]: https://github.com/AntixK/PyTorch-VAE/blob/master/assets/recons_SWAE_49.png

[python-image]: https://img.shields.io/badge/Python-3.5-ff69b4.svg
[python-url]: https://www.python.org/
Expand Down
4 changes: 3 additions & 1 deletion configs/swae.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,10 @@ model_params:
name: 'SWAE'
in_channels: 3
latent_dim: 128
reg_weight: 100
reg_weight: 10
wasserstein_deg: 2.0
num_projections: 100
projection_dist: "cauchy"

exp_params:
dataset: celeba
Expand Down
1 change: 0 additions & 1 deletion experiment.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,6 @@ def data_transforms(self):
transforms.CenterCrop(148),
transforms.Resize(self.params['img_size']),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
SetRange])
else:
raise ValueError('Undefined dataset type')
Expand Down
99 changes: 15 additions & 84 deletions models/swae.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
from models import BaseVAE
from torch import nn
from torch.nn import functional as F
from torch import distributions as dist
from .types_ import *


Expand All @@ -13,12 +14,16 @@ def __init__(self,
hidden_dims: List = None,
reg_weight: int = 100,
wasserstein_deg: float= 2.,
num_projections: int = 50,
projection_dist: str = 'normal',
**kwargs) -> None:
super(SWAE, self).__init__()

self.latent_dim = latent_dim
self.reg_weight = reg_weight
self.p = wasserstein_deg
self.num_projections = num_projections
self.proj_dist = projection_dist

modules = []
if hidden_dims is None:
Expand Down Expand Up @@ -121,74 +126,6 @@ def loss_function(self,
loss = recons_loss_l2 + recons_loss_l1 + swd_loss
return {'loss': loss, 'Reconstruction_Loss':(recons_loss_l2 + recons_loss_l1), 'SWD': swd_loss}

def compute_kernel(self,
x1: Tensor,
x2: Tensor) -> Tensor:
# Convert the tensors into row and column vectors
D = x1.size(1)
N = x1.size(0)

x1 = x1.unsqueeze(-2) # Make it into a column tensor
x2 = x2.unsqueeze(-3) # Make it into a row tensor

"""
Usually the below lines are not required, especially in our case,
but this is useful when x1 and x2 have different sizes
along the 0th dimension.
"""
x1 = x1.expand(N, N, D)
x2 = x2.expand(N, N, D)

if self.kernel_type == 'rbf':
result = self.compute_rbf(x1, x2)
elif self.kernel_type == 'imq':
result = self.compute_inv_mult_quad(x1, x2)
else:
raise ValueError('Undefined kernel type.')

return result


def compute_rbf(self,
x1: Tensor,
x2: Tensor,
eps: float = 1e-7) -> Tensor:
"""
Computes the RBF Kernel between x1 and x2.
:param x1: (Tensor)
:param x2: (Tensor)
:param eps: (Float)
:return:
"""
z_dim = x2.size(-1)
sigma = 2. * z_dim * self.z_var

result = torch.exp(-((x1 - x2).pow(2).mean(-1) / sigma))
return result

def compute_inv_mult_quad(self,
x1: Tensor,
x2: Tensor,
eps: float = 1e-7) -> Tensor:
"""
Computes the Inverse Multi-Quadratics Kernel between x1 and x2,
given by
k(x_1, x_2) = \sum \frac{C}{C + \|x_1 - x_2 \|^2}
:param x1: (Tensor)
:param x2: (Tensor)
:param eps: (Float)
:return:
"""
z_dim = x2.size(-1)
C = 2 * z_dim * self.z_var
kernel = C / (eps + C + (x1 - x2).pow(2).sum(dim = -1))

# Exclude diagonal elements
result = kernel.sum() - kernel.diag().sum()

return result

def get_random_projections(self, latent_dim: int, num_samples: int) -> Tensor:
"""
Returns random samples from latent distribution's (Gaussian)
Expand All @@ -199,7 +136,14 @@ def get_random_projections(self, latent_dim: int, num_samples: int) -> Tensor:
:param num_samples: (Int) Number of samples required (S)
:return: Random projections from the latent unit sphere
"""
rand_samples = torch.randn(num_samples, latent_dim)
if self.proj_dist == 'normal':
rand_samples = torch.randn(num_samples, latent_dim)
elif self.proj_dist == 'cauchy':
rand_samples = dist.Cauchy(torch.tensor([0.0]),
torch.tensor([1.0])).sample((num_samples, latent_dim)).squeeze()
else:
raise ValueError('Unknown projection distribution.')

rand_proj = rand_samples / rand_samples.norm(dim=1).view(-1,1)
return rand_proj # [S x D]

Expand All @@ -222,7 +166,7 @@ def compute_swd(self,
device = z.device

proj_matrix = self.get_random_projections(self.latent_dim,
num_samples=50).transpose(0,1).to(device)
num_samples=self.num_projections).transpose(0,1).to(device)

latent_projections = z.matmul(proj_matrix) # [N x S]
prior_projections = prior_z.matmul(proj_matrix) # [N x S]
Expand All @@ -234,19 +178,6 @@ def compute_swd(self,
w_dist = w_dist.pow(p)
return reg_weight * w_dist.mean()

def compute_mmd(self, z: Tensor, reg_weight: float) -> Tensor:
# Sample from prior (Gaussian) distribution
prior_z = torch.randn_like(z)

prior_z__kernel = self.compute_kernel(prior_z, prior_z)
z__kernel = self.compute_kernel(z, z)
priorz_z__kernel = self.compute_kernel(prior_z, z)

mmd = reg_weight * prior_z__kernel.mean() + \
reg_weight * z__kernel.mean() - \
2 * reg_weight * priorz_z__kernel.mean()
return mmd

def sample(self,
num_samples:int,
current_device: int, **kwargs) -> Tensor:
Expand All @@ -272,4 +203,4 @@ def generate(self, x: Tensor, **kwargs) -> Tensor:
:return: (Tensor) [B x C x H x W]
"""

return self.forward(x)[0]
return self.forward(x)[0]
7 changes: 1 addition & 6 deletions models/vanilla_vae.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -157,12 +157,7 @@ def sample(self,
z = torch.randn(num_samples,
self.latent_dim)

z = z.to(current_device) #
# vutils.save_image(test_input.data,
# f"{self.logger.save_dir}{self.logger.name}/version_{self.logger.version}/"
# f"real_img_{self.logger.name}_{self.current_epoch}.png",
# normalize=True,
# nrow=int(math.sqrt(self.params['batch_size'])))
z = z.to(current_device)

samples = self.decode(z)
return samples
Expand Down

0 comments on commit 76b4a6f

Please sign in to comment.