options, and bugfixes

readme.md

@@ -4,7 +4,7 @@ This project uses [Attentive Neural Process](https://arxiv.org/abs/1901.05761) (
 
 ![](docs/anp.png)
 
-This repository also includes a pytorch implementation that is more flexible and stable than others available now (as of 2019-11-01).
+This repository also includes a pytorch implementation that has been tweaked to be more flexible and stable. It may be usefull if you are looking for a ANP model in pytorch, and seems more stable than others available now (as of 2019-11-01).
 
 ## Usage
 
@@ -14,6 +14,12 @@ This repository also includes a pytorch implementation that is more flexible and
 
 ## Data
 - Some data is included, you can get more from https://www.kaggle.com/jeanmidev/smart-meters-in-london/version/11
+- Inputs are: 
+  - Weather
+  - Time features: time of day, day of week, month of year, etc
+  - Bank holidays
+  - Position in sequence: days since start of window
+- Target is: mean power usage on block
 
 ## Example outputs
 
@@ -30,22 +36,28 @@ Here the black dots are input data, the dotted line is the true data. The blue l
 
 ## Code
 
-This is based on the code listed in the next section, with some changes. The most notable ones add stability:
+This is based on the code listed in the next section, with some changes. The most notable ones add stability, others are to make sure it can handle predicting into the future:
 
 Changes for stability:
-- in eval mode, take mean of latent space, and mean output, don't sample
-- use log_variance where possible
+- in eval mode, take mean of latent space, and mean of output isntead of sampling
+- use log_variance where possible (there is a flag to try without this)
   - and add a minimum bound to std (in log domain) to avoid mode collapse
-- use pytorch attention (which has dropout)
+- use pytorch attention (which has dropout) instead of custom attention
 - use batchnorm and dropout on channel dimensions
-- use log_prob loss
+- use log_prob loss (no mseloss or BCELoss)
 - check and skip nonfinite values because for extreme inputs we can still get nan's
 
+Changes for a predictive use case:
+- target points are always in the future, context is in the past
+- context and and targets are still sampled randomly during training
+
 
 ## See also:
 
-- Original code in tensorflow: https://github.com/deepmind/neural-processes/blob/master/attentive_neural_process.ipynb
-- First pytorch implementation: https://github.com/soobinseo/Attentive-Neural-Process/blob/master/network.py
-- Second pytorch implementation (has some major bugs) https://github.com/KurochkinAlexey/Attentive-neural-processes/blob/master/anp_1d_regression.ipynb
+A list of projects I used as reference, is modified to make this one:
+
+- Original code in tensorflow from hyunjik11 (author of the original paper) : https://github.com/deepmind/neural-processes/blob/master/attentive_neural_process.ipynb
+- First pytorch implementation by soobinseo: https://github.com/soobinseo/Attentive-Neural-Process/blob/master/network.py
+- Second pytorch implementation KurochkinAlexey (has some bugs currently) https://github.com/KurochkinAlexey/Attentive-neural-processes/blob/master/anp_1d_regression.ipynb
 - If you want to try vanilla neural processes: https://github.com/EmilienDupont/neural-processes/blob/master/example-1d.ipynb
 

src/models/model.py

@@ -49,9 +49,15 @@ class LatentModel(nn.Module):
         n_decoder_layers=3,
         num_heads=8,
         dropout=0,
+        attention_dropout=0,
+        min_std=0.1,
+        use_lvar=True,
+        use_deterministic_path=True
     ):
 
         super().__init__()
+        self.use_lvar = use_lvar
+        self.use_deterministic_path = use_deterministic_path
 
         self._latent_encoder = LatentEncoder(
             x_dim + y_dim,
@@ -60,7 +66,10 @@ class LatentModel(nn.Module):
             self_attention_type=latent_enc_self_attn_type,
             n_encoder_layers=n_latent_encoder_layers,
             dropout=dropout,
+            attention_dropout=attention_dropout,
             n_heads=num_heads,
+            min_std=min_std,
+            use_lvar=use_lvar
         )
 
         self._deterministic_encoder = DeterministicEncoder(
@@ -71,6 +80,7 @@ class LatentModel(nn.Module):
             cross_attention_type=det_enc_cross_attn_type,
             n_d_encoder_layers=n_det_encoder_layers,
             dropout=dropout,
+            attention_dropout=attention_dropout,
             n_heads=num_heads,
         )
 
@@ -81,6 +91,8 @@ class LatentModel(nn.Module):
             latent_dim=latent_dim,
             n_decoder_layers=n_decoder_layers,
             dropout=dropout,
+            min_std=min_std,
+            use_lvar=use_lvar
         )
 
     def forward(self, context_x, context_y, target_x, target_y=None):
@@ -93,24 +105,32 @@ class LatentModel(nn.Module):
             if self.training:
                 z = dist_post.rsample()
             else:
+                # instead of sampling, in test mode take the mean, this will make it more deterministic
                 z = dist_post.loc
         else:
-            z = (
-                dist_prior.loc
-            )  # instead of sampling, in test mode take the mean, this will make it more deterministic
+            z = dist_prior.loc
 
         z = z.unsqueeze(1).repeat(1, num_targets, 1)  # [B, T_target, H]
-        r = self._deterministic_encoder(
-            context_x, context_y, target_x
-        )  # [B, T_target, H]
+
+        if self.use_deterministic_path:
+            r = self._deterministic_encoder(
+                context_x, context_y, target_x
+            )  # [B, T_target, H]
+        else:
+            r = None
         dist, log_sigma = self._decoder(r, z, target_x)
+
         if target_y is not None:
-            # Log likelihood has shape (batch_size, num_target, y_dim).
-            log_p = log_prob_sigma(target_y, dist.loc, log_sigma).mean(-1)
-            #  KL has shape (batch_size, r_dim)
-            kl_loss = kl_loss_var(
-                dist_prior.loc, log_var_prior, dist_post.loc, log_var_post
-            ).mean(-1)
+            if self.use_lvar:
+                # Log likelihood has shape (batch_size, num_target, y_dim).
+                log_p = log_prob_sigma(target_y, dist.loc, log_sigma).mean(-1)
+                #  KL has shape (batch_size, r_dim)
+                kl_loss = kl_loss_var(
+                    dist_prior.loc, log_var_prior, dist_post.loc, log_var_post
+                ).mean(-1)
+            else:
+                log_p = dist.log_prob(target_y).mean(-1)
+                kl_loss = torch.distributions.kl_divergence(dist_post, dist_prior).mean(-1)
             kl_loss = kl_loss[:, None].expand(log_p.shape)
             loss = (kl_loss - log_p).mean()
 

src/models/modules.py

@@ -90,6 +90,9 @@ class Attention(nn.Module):
 
 
 class LatentEncoder(nn.Module):
+    """
+    Latent Encoder [For prior, posterior]
+    """
     def __init__(
         self,
         input_dim,
@@ -100,8 +103,11 @@ class LatentEncoder(nn.Module):
         n_encoder_layers=3,
         min_std=0.1,
         dropout=0,
+        attention_dropout=0,
+        use_lvar=False,
     ):
         super().__init__()
+        self.use_lvar = use_lvar
         self._input_layer = NPBlockRelu2d(input_dim, hidden_dim, dropout)
         self._encoder = nn.Sequential(
             *[
@@ -110,7 +116,7 @@ class LatentEncoder(nn.Module):
             ]
         )
         self._self_attention = Attention(
-            hidden_dim, self_attention_type, n_heads=n_heads, dropout=dropout
+            hidden_dim, self_attention_type, n_heads=n_heads, dropout=attention_dropout
         )
         self._penultimate_layer = block_relu(hidden_dim, hidden_dim, dropout)
         self._mean = nn.Linear(hidden_dim, latent_dim)
@@ -118,29 +124,51 @@ class LatentEncoder(nn.Module):
         self.min_std = min_std
 
     def forward(self, x, y):
-        encoder_input = torch.cat([x, y], dim=-1)
-        encoded = self._input_layer(encoder_input)
+        """Encodes the inputs into one representation.
 
+        Args:
+        x: Tensor of shape [B,observations,d_x]. For this 1D regression
+            task this corresponds to the x-values.
+        y: Tensor of shape [B,observations,d_y]. For this 1D regression
+            task this corresponds to the y-values.
+
+        Returns:
+        - A normal distribution over tensors of shape [B, num_latents]
+        - log_var
+        """
+        # Concat location (x) and value (y) along the filter axes
+        encoder_input = torch.cat([x, y], dim=-1)
+
+        # Pass final axis through MLP
+        encoded = self._input_layer(encoder_input)
         encoded = self._encoder(encoded)
 
+        # Self-attention aggregator
         attention_output = self._self_attention(encoded, encoded, encoded)
-
         mean_repr = attention_output.mean(dim=1)
 
-        mean_repr = torch.relu(self._penultimate_layer(mean_repr))
+        # Have further MLP layers that map to the parameters of the Gaussian latent
+        mean_repr = self._penultimate_layer(mean_repr)
 
+        # Then apply further linear layers to output latent mu and log sigma
         mean = self._mean(mean_repr)
         log_var = self._log_var(mean_repr)
 
-        # Clip it in the log domain, so it can only approach self.min_std, this helps aboid mode collaprse
-        log_var = F.softplus(log_var) + math.log(self.min_std)
-
-        sigma = torch.exp(0.5 * log_var)
+        # Clip it in the log domain, so it can only approach self.min_std, this helps aboid mode collapase
+        if self.use_lvar:
+            log_var = log_var + math.log(self.min_std)
+            sigma = torch.exp(0.5 * log_var)
+        else:
+            sigma = self.min_std + (1 - self.min_std) * torch.sigmoid(log_var * 0.5)
         dist = torch.distributions.Normal(mean, sigma)
         return dist, log_var
 
 
 class DeterministicEncoder(nn.Module):
+    """
+    Deterministic Encoder [r]
+    """
+
     def __init__(
         self,
         input_dim,
@@ -150,6 +178,7 @@ class DeterministicEncoder(nn.Module):
         self_attention_type="multihead",
         cross_attention_type="multihead",
         dropout=0,
+        attention_dropout=0,
         n_heads=8,
     ):
         super().__init__()
@@ -161,25 +190,32 @@ class DeterministicEncoder(nn.Module):
             ]
         )
         self._self_attention = Attention(
-            hidden_dim, self_attention_type, dropout=dropout, n_heads=n_heads
+            hidden_dim, self_attention_type, dropout=attention_dropout, n_heads=n_heads
         )
         self._cross_attention = Attention(
-            hidden_dim, cross_attention_type, dropout=dropout, n_heads=n_heads
+            hidden_dim, cross_attention_type, dropout=attention_dropout, n_heads=n_heads
         )
         self._target_transform = nn.Linear(x_dim, hidden_dim)
         self._context_transform = nn.Linear(x_dim, hidden_dim)
 
     def forward(self, context_x, context_y, target_x):
+        # concat context location (x), context value (y)
         d_encoder_input = torch.cat([context_x, context_y], dim=-1)
+
+        # Pass final axis through MLP
         d_encoded = self._input_layer(d_encoder_input)
         d_encoded = self._d_encoder(d_encoded)
-        attention_output = self._self_attention(d_encoded, d_encoded, d_encoded)
 
-        q = self._target_transform(target_x)
+        # Apply self attention
+        d_encoded = self._self_attention(d_encoded, d_encoded, d_encoded)
+
+        # query: target_x, key: context_x, value: d_encoded (representation of x)
         k = self._context_transform(context_x)
-        q = self._cross_attention(k, d_encoded, q)
+        q = self._target_transform(target_x)
 
-        return q
+        # Cross Attention
+        r = self._cross_attention(k, d_encoded, q)
+        return r
 
 
 class Decoder(nn.Module):
@@ -192,10 +228,12 @@ class Decoder(nn.Module):
         n_decoder_layers=3,
         min_std=0.1,
         dropout=0,
+        use_lvar=False,
     ):
         super().__init__()
+        self.use_lvar = use_lvar
         self._target_transform = NPBlockRelu2d(x_dim, hidden_dim, dropout)
-        hidden_dim_2 = 2* hidden_dim + latent_dim
+        hidden_dim_2 = 2 * hidden_dim + latent_dim
         self._decoder = nn.Sequential(
             *[
                 NPBlockRelu2d(hidden_dim_2, hidden_dim_2, dropout)
@@ -208,14 +246,26 @@ class Decoder(nn.Module):
 
     def forward(self, r, z, target_x):
         x = self._target_transform(target_x)
-        representation = torch.cat([torch.cat([r, z], dim=-1), x], dim=-1)
+
+        # concatenate target_x and representation
+        if r is not None:
+            z = torch.cat([r, z], dim=-1)
+        representation = torch.cat([z, x], dim=-1)
+
+        # Pass final axis through MLP
         representation = self._decoder(representation)
 
+        # Get the mean and the variance
         mean = self._mean(representation)
         log_sigma = self._std(representation)
 
-        log_sigma = F.softplus(log_sigma) + math.log(self.min_std)
-        sigma = torch.exp(log_sigma)
+        # Bound the variance
+        if self.use_lvar:
+            log_sigma = log_sigma + math.log(self.min_std)
+            sigma = torch.exp(log_sigma)
+        else:
+            sigma = self.min_std + (1-self.min_std) * F.softplus(log_sigma)
 
+        # Dist
         dist = torch.distributions.Normal(mean, sigma)
         return dist, log_sigma