Add 'CVAE' model

Author: ffiirree

README.md

@@ -98,6 +98,8 @@
 ### VAEs
 
 - [x] [`VAE`](cvm/models/vae/vae.py) - [Auto-Encoding Variational Bayes](https://arxiv.org/abs/1312.6114), 2013
+- [x] [`CVAE`](cvm/models/vae/cvae.py) - [Learning Structured Output Representation using Deep Conditional Generative Models
+](https://papers.nips.cc/paper/2015/hash/8d55a249e6baa5c06772297520da2051-Abstract.html), NeurIPS, 2015
 - [ ] `β-VAE` - [beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework](https://openreview.net/forum?id=Sy2fzU9gl), ICLR, 2017
 
 

cvm/models/vae/__init__.py

@@ -1 +1,2 @@
-from .vae import *
+from .vae import *
+from .cvae import *

cvm/models/vae/cvae.py

@@ -0,0 +1,81 @@
+import torch
+import torch.nn as nn
+
+from ..utils import export, load_from_local_or_url
+from typing import Any
+
+
+@export
+class ConditionalVAE(nn.Module):
+    """
+        Paper: [Learning Structured Output Representation using Deep Conditional Generative Models](https://papers.nips.cc/paper/2015/hash/8d55a249e6baa5c06772297520da2051-Abstract.html)
+    """
+    def __init__(
+        self,
+        image_size,
+        nz: int = 100,
+        **kwargs: Any
+    ):
+        super().__init__()
+
+        self.image_size = image_size
+        self.nz = nz
+
+        self.embeds_en = nn.Embedding(10, 200)
+
+        self.embeds_de = nn.Embedding(10, 10)
+
+        # Q(z|X)
+        self.encoder = nn.Sequential(
+            nn.Linear(self.image_size ** 2 + 200, 512),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(512, 512),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(512, 256),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(256, self.nz * 2)
+        )
+
+        # P(X|z)
+        self.decoder = nn.Sequential(
+            nn.Linear(self.nz + 10, 256),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(256, 512),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(512, 512),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(512, self.image_size ** 2),
+            nn.Sigmoid(),
+            nn.Unflatten(1, (1, image_size, image_size))
+        )
+
+    def sample_z(self, mu, logvar, c):
+        eps = torch.randn_like(logvar)
+
+        return torch.cat([mu + eps * torch.exp(0.5 * logvar), c], dim=1)
+
+    def forward(self, x, c):
+        x = torch.flatten(x, 1)
+
+        x = torch.cat([x, self.embeds_en(c)], dim=1)
+
+        mu, logvar = torch.chunk(self.encoder(x), 2, dim=1)
+
+        z = self.sample_z(mu, logvar, self.embeds_de(c))
+
+        x = self.decoder(z)
+        return x, mu, logvar
+
+
+@export
+def cvae(
+    pretrained: bool = False,
+    pth: str = None,
+    progress: bool = True,
+    **kwargs: Any
+):
+    model = ConditionalVAE(**kwargs)
+
+    if pretrained:
+        load_from_local_or_url(model, pth, kwargs.get('url', None), progress)
+    return model

cvm/version.py

@@ -1 +1 @@
-__version__ = '0.1.20'
+__version__ = '0.1.21'