Merge pull request #9 from ziatdinovmax/master

Add joint (discrete + continuous) VAE

.gitpod.yml

@@ -1,6 +1,2 @@
 image:
   file: .gitpod.Dockerfile 
-ports:
- - port: 8080
-tasks:
- - command: jupyter lab --no-browser --port 8080

atomai/__version__.py

@@ -1 +1 @@
-version = '0.6.2'
+version = '0.6.5'

atomai/losses_metrics/__init__.py

@@ -1,4 +1,9 @@
-from .losses import dice_loss, focal_loss, select_loss, vae_loss, rvae_loss
+from .losses import dice_loss, focal_loss, select_loss
 from .metrics import IoU
+from .vi_losses import (joint_rvae_loss, joint_vae_loss, kld_discrete,
+                        kld_normal, kld_rot, rvae_loss, vae_loss,
+                        reconstruction_loss)
 
-__all__ = ['focal_loss', 'dice_loss', 'select_seg_loss', 'IoU']
+__all__ = ['focal_loss', 'dice_loss', 'select_loss', "vae_loss",
+           "rvae_loss", "joint_vae_loss", "joint_rvae_loss", "IoU",
+           "kld_normal", "kld_discrete", "kld_rot", "reconstruction_loss"]

atomai/losses_metrics/losses.py

@@ -4,8 +4,6 @@
 
 Custom Pytorch loss functions
 """
-from typing import Tuple
-import numpy as np
 import torch
 import torch.nn.functional as F
 
@@ -89,79 +87,9 @@ def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
         return (1 - dice_loss)
 
 
-def vae_loss(reconstr_loss: str,
-             in_dim: Tuple[int],
-             x: torch.Tensor,
-             x_reconstr: torch.Tensor,
-             *args: torch.Tensor
-             ) -> torch.Tensor:
-    """
-    Calculates ELBO
-    """
-    batch_dim = x.size(0)
-    if len(args) == 2:
-        z_mean, z_logsd = args
-    else:
-        z_mean = z_logsd = torch.zeros((batch_dim, 1))
-    z_sd = torch.exp(z_logsd)
-    if reconstr_loss == "mse":
-        reconstr_error = -0.5 * torch.sum(
-            (x_reconstr.reshape(batch_dim, -1) - x.reshape(batch_dim, -1))**2, 1).mean()
-    elif reconstr_loss == "ce":
-        px_size = np.product(in_dim)
-        rs = (np.product(in_dim[:2]),)
-        if len(in_dim) == 3:
-            rs = rs + (in_dim[-1],)
-        reconstr_error = -F.binary_cross_entropy_with_logits(
-            x_reconstr.reshape(-1, *rs), x.reshape(-1, *rs)) * px_size
-    else:
-        raise NotImplementedError("Reconstruction loss must be 'mse' or 'ce'")
-    kl_z = -z_logsd + 0.5 * z_sd**2 + 0.5 * z_mean**2 - 0.5
-    kl_z = torch.sum(kl_z, 1).mean()
-    return reconstr_error - kl_z
-
-
-def rvae_loss(reconstr_loss: str,
-              in_dim: Tuple[int],
-              x: torch.Tensor,
-              x_reconstr: torch.Tensor,
-              *args: torch.Tensor,
-              **kwargs: float) -> torch.Tensor:
-    """
-    Calculates ELBO
-    """
-    batch_dim = x.size(0)
-    if len(args) == 2:
-        z_mean, z_logsd = args
-    else:
-        z_mean = z_logsd = torch.zeros((batch_dim, 1))
-    phi_prior = kwargs.get("phi_prior", 0.1)
-    z_sd = torch.exp(z_logsd)
-    phi_sd, phi_logsd = z_sd[:, 0], z_logsd[:, 0]
-    z_mean, z_sd, z_logsd = z_mean[:, 1:], z_sd[:, 1:], z_logsd[:, 1:]
-    batch_dim = x.size(0)
-    if reconstr_loss == "mse":
-        reconstr_error = -0.5 * torch.sum(
-            (x_reconstr.view(batch_dim, -1) - x.view(batch_dim, -1))**2, 1).mean()
-    elif reconstr_loss == "ce":
-        px_size = np.product(in_dim)
-        rs = (np.product(in_dim[:2]),)
-        if len(in_dim) == 3:
-            rs = rs + (in_dim[-1],)
-        reconstr_error = -F.binary_cross_entropy_with_logits(
-            x_reconstr.view(-1, *rs), x.view(-1, *rs)) * px_size
-    else:
-        raise NotImplementedError("Reconstruction loss must be 'mse' or 'ce'")
-    kl_rot = (-phi_logsd + np.log(phi_prior) +
-                phi_sd**2 / (2 * phi_prior**2) - 0.5)
-    kl_z = -z_logsd + 0.5 * z_sd**2 + 0.5 * z_mean**2 - 0.5
-    kl_div = (kl_rot + torch.sum(kl_z, 1)).mean()
-    return reconstr_error - kl_div
-
-
 def select_loss(loss: str, nb_classes: int = None):
     """
-    Selects loss for a semantic segmentation model training
+    Selects loss for DCNN model training
     """
     if loss == 'ce' and nb_classes is None:
         raise ValueError("For cross-entropy loss function, you must" +

atomai/losses_metrics/vi_losses.py

@@ -0,0 +1,254 @@
+"""
+vae_losses.py
+=========
+
+Custom loss functions for Variational Autoencoders (VAEs)
+"""
+from typing import Tuple, List, Union, Optional
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+def reconstruction_loss(loss_type: str,
+                        in_dim: Tuple[int],
+                        x: torch.Tensor,
+                        x_reconstr: torch.Tensor,
+                        logits: bool = True,
+                        ) -> torch.Tensor:
+    """
+    Computes reconstruction loss (mse or cross-entropy)
+    without mean reduction (used in VAE objectives)
+    """
+    batch_dim = x.size(0)
+    if loss_type == "mse":
+        reconstr_loss = 0.5 * torch.sum(
+            (x_reconstr.reshape(batch_dim, -1) - x.reshape(batch_dim, -1))**2, 1)
+    elif loss_type == "ce":
+        rs = (np.product(in_dim[:2]),)
+        if len(in_dim) == 3:
+            rs = rs + (in_dim[-1],)
+        xe = (F.binary_cross_entropy_with_logits if
+              logits else F.binary_cross_entropy)
+        reconstr_loss = xe(x_reconstr.reshape(-1, *rs), x.reshape(-1, *rs),
+                           reduction='none').sum(-1)
+    else:
+        raise NotImplementedError("Reconstruction loss must be 'mse' or 'ce'")
+    return reconstr_loss
+
+
+def kld_normal(q_param: Tuple[torch.Tensor],
+               p_param: Optional[Tuple[torch.Tensor]] = None
+               ) -> torch.Tensor:
+    """
+    Kullback–Leibler (KL) divergence between two normal distributions
+    """
+    mu_1, log_sd_1 = q_param
+    sd_1 = torch.exp(log_sd_1)
+    if p_param is None:
+        # KL divergence b/w normal and standard normal distributions
+        kl = -log_sd_1 + 0.5 * sd_1**2 + 0.5 * mu_1**2 - 0.5
+    else:
+        mu_2, log_sd_2 = p_param
+        sd_2 = torch.exp(log_sd_2)
+        # KL divergence b/w two normal distributions
+        kl = (log_sd_2 - log_sd_1 +
+              0.5 * (sd_1**2 + (mu_1 - mu_2)**2) / sd_2**2 - 0.5)
+    return torch.sum(kl, -1)
+
+
+def kld_discrete(alpha: torch.Tensor):
+    """
+    Calculates the KL divergence between a Gumbel-Softmax distribution
+    and a uniform categorical distribution.
+
+    Args:
+        alpha:
+            Parameters of the Gumbel-Softmax distribution.
+    """
+    eps = 1e-12
+    cat_dim = alpha.size(-1)
+    h1 = torch.log(alpha + eps)
+    h2 = np.log(1. / cat_dim + eps)
+    kld_loss = torch.mean(torch.sum(alpha * (h1 - h2), dim=1), dim=0)
+    return kld_loss.view(1)
+
+
+def kld_rot(phi_prior: torch.Tensor, phi_logsd: torch.Tensor) -> torch.Tensor:
+    """
+    Kullback–Leibler (KL) divergence for rotation latent variable
+    """
+    phi_sd = torch.exp(phi_logsd)
+    kl_rot = (-phi_logsd + np.log(phi_prior) +
+              phi_sd**2 / (2 * phi_prior**2) - 0.5)
+    return kl_rot
+
+
+def vae_loss(recon_loss: str,
+             in_dim: Tuple[int],
+             x: torch.Tensor,
+             x_reconstr: torch.Tensor,
+             *args: torch.Tensor,
+             ) -> torch.Tensor:
+    """
+    Calculates ELBO
+    """
+    if len(args) == 2:
+        q_param = args
+    else:
+        raise ValueError(
+            "Pass mean and SD values of encoded distribution as args")
+    likelihood = -reconstruction_loss(recon_loss, in_dim, x, x_reconstr).mean()
+    kl_z = kld_normal(q_param).mean()
+    return likelihood - kl_z
+
+
+def rvae_loss(recon_loss: str,
+              in_dim: Tuple[int],
+              x: torch.Tensor,
+              x_reconstr: torch.Tensor,
+              *args: torch.Tensor,
+              **kwargs: float) -> torch.Tensor:
+    """
+    Calculates ELBO
+    """
+    if len(args) == 2:
+        z_mean, z_logsd = args
+    else:
+        raise ValueError(
+            "Pass mean and SD values of encoded distribution as args")
+    phi_prior = kwargs.get("phi_prior", 0.1)
+    b1, b2 = kwargs.get("b1", 1), kwargs.get("b2", 1)
+    phi_logsd = z_logsd[:, 0]
+    z_mean, z_logsd = z_mean[:, 1:], z_logsd[:, 1:]
+    likelihood = -reconstruction_loss(recon_loss, in_dim, x, x_reconstr).mean()
+    kl_rot = kld_rot(phi_prior, phi_logsd).mean()
+    kl_z = kld_normal([z_mean, z_logsd]).mean()
+    kl_div = (b1*kl_z + b2 * kl_rot)
+    return likelihood - kl_div
+
+
+def joint_vae_loss(recon_loss: str,
+                   in_dim: Tuple[int],
+                   x: torch.Tensor,
+                   x_reconstr: torch.Tensor,
+                   *args: torch.Tensor,
+                   **kwargs: Union[List, int],
+                   ) -> torch.Tensor:
+    """
+    Calculates joint ELBO for continuous and discrete variables
+    """
+    if len(args) == 3:
+        z_mean, z_logsd, alphas = args
+    else:
+        raise ValueError(
+            "Pass continuous (mean, SD) and discrete (alphas) values" +
+            "of encoded distributions as args")
+
+    cont_capacity = kwargs.get("cont_capacity", [0.0, 5.0, 25000, 30])
+    disc_capacity = kwargs.get("disc_capacity", [0.0, 5.0, 25000, 30])
+    num_iter = kwargs.get("num_iter", 0)
+    disc_dims = [a.size(1) for a in alphas]
+
+    # Calculate reconstruction loss term
+    likelihood = -reconstruction_loss(recon_loss, in_dim, x, x_reconstr).mean()
+
+    # Calculate KL term for continuous latent variables
+    kl_cont_loss = kld_normal([z_mean, z_logsd]).mean()
+    # Calculate KL term for discrete latent variables
+    kl_disc = [kld_discrete(alpha) for alpha in alphas]
+    kl_disc_loss = torch.sum(torch.cat(kl_disc))
+
+    # Apply information capacity terms to contninuous and discrete channels
+    cargs = [kl_cont_loss, kl_disc_loss, cont_capacity,
+             disc_capacity, disc_dims, num_iter]
+    cont_capacity_loss, disc_capacity_loss = infocapacity(*cargs)
+
+    return likelihood - cont_capacity_loss - disc_capacity_loss
+
+
+def joint_rvae_loss(recon_loss: str,
+                    in_dim: Tuple[int],
+                    x: torch.Tensor,
+                    x_reconstr: torch.Tensor,
+                    *args: torch.Tensor,
+                    **kwargs: float) -> torch.Tensor:
+    """
+    Calculates joint ELBO for continuous and discrete variables
+    """
+    if len(args) == 3:
+        z_mean, z_logsd, alphas = args
+    else:
+        raise ValueError(
+            "Pass continuous (mean, SD) and discrete (alphas) values" +
+            "of encoded distributions as args")
+
+    phi_prior = kwargs.get("phi_prior", 0.1)
+    klrot_cap = kwargs.get("klrot_cap", True)
+    cont_capacity = kwargs.get("cont_capacity", [0.0, 5.0, 25000, 30])
+    disc_capacity = kwargs.get("disc_capacity", [0.0, 5.0, 25000, 30])
+    num_iter = kwargs.get("num_iter", 0)
+
+    # Calculate reconstruction loss term
+    likelihood = -reconstruction_loss(recon_loss, in_dim, x, x_reconstr).mean()
+
+    # Calculate KL term for continuous latent variables
+    phi_logsd = z_logsd[:, 0]  # rotation
+    z_mean, z_logsd = z_mean[:, 1:], z_logsd[:, 1:]  # image content
+    kl_rot = kld_rot(phi_prior, phi_logsd).mean()
+    kl_z = kld_normal([z_mean, z_logsd]).mean()
+    if klrot_cap:
+        kl_cont_loss = kl_z + kl_rot
+    else:  # no capacity limit on KL term associated with rotations
+        kl_cont_loss = kl_z
+
+    # Calculate KL term for discrete latent variables
+    disc_dims = [a.size(1) for a in alphas]
+    kl_disc = [kld_discrete(alpha) for alpha in alphas]
+    kl_disc_loss = torch.sum(torch.cat(kl_disc))
+
+    # Apply information capacity terms to contninuous and discrete channels
+    cargs = [kl_cont_loss, kl_disc_loss, cont_capacity,
+             disc_capacity, disc_dims, num_iter]
+    cont_capacity_loss, disc_capacity_loss = infocapacity(*cargs)
+    if not klrot_cap:
+        cont_capacity_loss = cont_capacity_loss + kl_rot
+
+    return likelihood - cont_capacity_loss - disc_capacity_loss
+
+
+def infocapacity(kl_cont_loss: torch.Tensor,
+                 kl_disc_loss: torch.Tensor,
+                 cont_capacity: List[float],
+                 disc_capacity: List[float],
+                 disc_dims: List[int],
+                 num_iter: int) -> torch.Tensor:
+    """
+    Controls information capacity of the continuous and discrete loss
+    (based on https://arxiv.org/pdf/1804.00104.pdf &
+    https://github.com/Schlumberger/joint-vae/blob/master/jointvae/training.py)
+    """
+    # Linearly increase capacity of continuous channels
+    cont_min, cont_max, cont_num_iters, cont_gamma = cont_capacity
+    # Increase continuous capacity without exceeding cont_max
+    cont_cap_current = (cont_max - cont_min) * num_iter
+    cont_cap_current = cont_cap_current / float(cont_num_iters) + cont_min
+    cont_cap_current = min(cont_cap_current, cont_max)
+    # Calculate continuous capacity loss
+    cont_capacity_loss = cont_gamma*torch.abs(cont_cap_current - kl_cont_loss)
+
+    # Linearly increase capacity of discrete channels
+    disc_min, disc_max, disc_num_iters, disc_gamma = disc_capacity
+    # Increase discrete capacity without exceeding disc_max or theoretical
+    # maximum (i.e. sum of log of dimension of each discrete variable)
+    disc_cap_current = (disc_max - disc_min) * num_iter
+    disc_cap_current = disc_cap_current / float(disc_num_iters) + disc_min
+    disc_cap_current = min(disc_cap_current, disc_max)
+    # Require float conversion here to not end up with numpy float
+    disc_theory_max = sum([float(np.log(d)) for d in disc_dims])
+    disc_cap_current = min(disc_cap_current, disc_theory_max)
+    # Calculate discrete capacity loss
+    disc_capacity_loss = disc_gamma*torch.abs(disc_cap_current - kl_disc_loss)
+
+    return cont_capacity_loss, disc_capacity_loss
+

atomai/models/__init__.py

@@ -1,7 +1,7 @@
 from .segmentor import Segmentor
 from .imspec import ImSpec
-from .vae import BaseVAE, VAE, rVAE
+from .dgm import BaseVAE, VAE, rVAE, jVAE, jrVAE
 from .loaders import load_model, load_ensemble
 
 __all__ = ["Segmentor", "ImSpec", "BaseVAE", "VAE", "rVAE",
-           "load_model", "load_ensemble"]
+           "jVAE", "jrVAE", "load_model", "load_ensemble"]

atomai/models/dgm/__init__.py

@@ -0,0 +1,6 @@
+from .vae import BaseVAE, VAE
+from .rvae import rVAE
+from .jvae import jVAE
+from .jrvae import jrVAE
+
+__all__ = ["BaseVAE", "VAE", "rVAE", "jVAE", "jrVAE"]