refactoring

anp-rnn_1d_regression.ipynb

@@ -99,8 +99,8 @@
    },
    "outputs": [],
    "source": [
-    "from src.data.gp_curves import GPCurvesReader\n",
-    "from src.models.model import LatentModel"
+    "from neural_processes.data.gp_curves import GPCurvesReader\n",
+    "from neural_processes.models.model import LatentModel"
    ]
   },
   {

neural_processes/lightning.py

@@ -0,0 +1,180 @@
+import pytorch_lightning as pl
+import torch
+import optuna
+import torch
+from torch import nn
+from pathlib import Path
+from pytorch_lightning.callbacks import EarlyStopping
+from optuna.integration.pytorch_lightning import _check_pytorch_lightning_availability
+import torch.nn.functional as F
+from .utils import ObjectDict
+from .data.smart_meter import get_smartmeter_df
+from .logger import logger
+
+
+class PL_Seq2Seq(pl.LightningModule):
+    def __init__(self, hparams, loss_fn=F.mse_loss, num_workers=3, MODEL_CLS=None):
+        super().__init__()
+        self.hparams = ObjectDict()
+        self.hparams.update(
+            hparams.__dict__ if hasattr(hparams, "__dict__") else hparams
+        )
+        self.num_workers = num_workers
+        self._model = MODEL_CLS(self.hparams)
+        self._datasets = None
+        self.loss_fn = loss_fn
+        self.train_logs = []  # HACK
+        self._dfs = None
+        # TODO make label name configurable
+        # TODO make data source configurable
+
+    def forward(self, *args, **kwargs):
+        return self._model(*args, **kwargs)
+
+    def training_step(self, batch, batch_idx):
+        # REQUIRED
+        assert all(torch.isfinite(d).all() for d in batch)
+        context_x, context_y, target_x, target_y = batch
+        y_dist, losses, extra = self.forward(context_x, context_y, target_x, target_y)
+        loss = losses['loss_p'] # + loss_mse
+        tensorboard_logs = {
+            "train/loss": loss,
+            'train/loss_mse': losses['loss_mse'],
+            "train/loss_p": losses['loss_p'],
+            "train/sigma": torch.exp(extra['log_sigma']).mean()}
+        return {"loss": loss, "log": tensorboard_logs}
+
+    def validation_step(self, batch, batch_idx):
+        context_x, context_y, target_x, target_y = batch
+        assert all(torch.isfinite(d).all() for d in batch)
+        y_dist, losses, extra = self.forward(context_x, context_y, target_x, target_y)
+        loss = losses['loss_p'] # + loss_mse
+        tensorboard_logs = {
+            "val_loss": loss,
+            'val/loss_mse': losses['loss_mse'],
+            "val/loss_p": losses['loss_p'],
+            "val/sigma": torch.exp(extra['log_sigma']).mean()}
+        return {"val_loss": loss, "log": tensorboard_logs}
+
+    def validation_end(self, outputs):
+        if int(self.hparams["vis_i"]) > 0:
+            self.show_image()
+
+        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
+        keys = outputs[0]["log"].keys()
+        tensorboard_logs = {
+            k: torch.stack([x["log"][k] for x in outputs if k in x["log"]]).mean()
+            for k in keys
+        }
+        tensorboard_logs_str = {k: f"{v}" for k, v in tensorboard_logs.items()}
+        print(f"step {self.trainer.global_step}, {tensorboard_logs_str}")
+        assert torch.isfinite(avg_loss)
+        return {"avg_val_loss": avg_loss, "log": tensorboard_logs}
+
+    def agg_logs(self, outputs):
+        if isinstance(outputs, dict):
+            outputs = [outputs]
+        aggs = {}
+        if len(outputs)>0:
+            for j in outputs[0]:
+                if isinstance(outputs[0][j], dict):
+                    # Take mean of sub dicts
+                    keys = outputs[0][j].keys()
+                    aggs[j] = {k: torch.stack([x[j][k] for x in outputs if k in x[j]]).mean().item() for k in keys}
+                else:
+                    # Take mean of numbers
+                    aggs[j] = torch.stack([x[j] for x in outputs if j in x]).mean().item()
+        return aggs
+
+    def show_image(self):        
+        # https://github.com/PytorchLightning/pytorch-lightning/blob/f8d9f8f/pytorch_lightning/core/lightning.py#L293
+        loader = self.val_dataloader()
+        vis_i = min(int(self.hparams["vis_i"]), len(loader.dataset))
+        # print('vis_i', vis_i)
+        if isinstance(self.hparams["vis_i"], str):
+            image = plot_from_loader(loader, self, i=int(vis_i))
+            plt.show()
+        else:
+            image = plot_from_loader_to_tensor(loader, self, i=vis_i)
+            self.logger.experiment.add_image('val/image', image, self.trainer.global_step)
+
+    def test_step(self, batch, batch_idx):
+        pred, losses, extra = self.forward(*batch)
+        # For test use a diff loss, MSE over next 24
+        # loss = losses["loss"]
+        loss = F.mse_loss(pred, batch[-1], reduction='none')[:, :24].mean()
+        tensorboard_logs = {"test_" + k: v for k, v in losses.items()}
+        return {"test_loss": loss, "log": tensorboard_logs}
+
+    def test_end(self, outputs):
+        avg_loss = torch.stack([x["test_loss"] for x in outputs]).mean()
+        keys = outputs[0]["log"].keys()
+        tensorboard_logs = {
+            k: torch.stack([x["log"][k] for x in outputs if k in x["log"]]).mean()
+            for k in keys
+        }
+        tensorboard_logs_str = {k: f"{v}" for k, v in tensorboard_logs.items()}
+
+        logger.info(
+            f"step {self.trainer.global_step}, {tensorboard_logs_str}"
+        )
+        return {"avg_val_loss": avg_loss, "log": tensorboard_logs}
+
+    def configure_optimizers(self):
+        optim = torch.optim.Adam(self.parameters(), lr=self.hparams["learning_rate"])
+        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+            optim, patience=self.hparams["patience"], verbose=True, min_lr=1e-7
+        )  # note early stopping has patience 3
+        return [optim], [scheduler]
+
+    def _get_cache_dfs(self):
+        if self._dfs is None:
+            df_train, df_val, df_test = get_smartmeter_df()
+            self._dfs = dict(df_train=df_train, df_val=df_val, df_test=df_test)
+        return self._dfs
+
+    @pl.data_loader
+    def train_dataloader(self):
+        df_train = self._get_cache_dfs()['df_train']
+        data_train = SmartMeterDataSet(
+            df_train, self.hparams["num_context"], self.hparams["num_extra_target"]
+        )
+        return torch.utils.data.DataLoader(
+            data_train,
+            batch_size=self.hparams["batch_size"],
+            shuffle=True,
+            collate_fn=collate_fns(
+                self.hparams["num_context"], self.hparams["num_extra_target"], sample=True, context_in_target=self.hparams["context_in_target"]
+            ),
+            num_workers=self.hparams["num_workers"],
+        )
+
+    @pl.data_loader
+    def val_dataloader(self):
+        df_test = self._get_cache_dfs()['df_val']
+        data_test = SmartMeterDataSet(
+            df_test, self.hparams["num_context"], self.hparams["num_extra_target"]
+        )
+        return torch.utils.data.DataLoader(
+            data_test,
+            batch_size=self.hparams["batch_size"],
+            shuffle=False,
+            collate_fn=collate_fns(
+                self.hparams["num_context"], self.hparams["num_extra_target"], sample=False, context_in_target=self.hparams["context_in_target"]
+            ),
+        )
+
+    @pl.data_loader
+    def test_dataloader(self):
+        df_test = self._get_cache_dfs()['df_test']
+        data_test = SmartMeterDataSet(
+            df_test, self.hparams["num_context"], self.hparams["num_extra_target"]
+        )
+        return torch.utils.data.DataLoader(
+            data_test,
+            batch_size=self.hparams["batch_size"],
+            shuffle=False,
+            collate_fn=collate_fns(
+                self.hparams["num_context"], self.hparams["num_extra_target"], sample=False, context_in_target=self.hparams["context_in_target"]
+            ),
+        )

neural_processes/logger.py

@@ -0,0 +1,4 @@
+import logging
+import sys
+logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+logger = logging.getLogger('predict_heading2')

neural_processes/models/lstm_seqseq.py

@@ -0,0 +1,147 @@
+import os
+import numpy as np
+import pandas as pd
+import torch
+from tqdm.auto import tqdm
+from torch import nn
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torchvision.datasets import MNIST
+from test_tube import Experiment, HyperOptArgumentParser
+from neural_processes.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df
+import torchvision.transforms as transforms
+from neural_processes.plot import plot_from_loader_to_tensor, plot_from_loader
+from argparse import ArgumentParser
+import json
+import pytorch_lightning as pl
+import math
+from matplotlib import pyplot as plt
+import torch
+import io
+import PIL
+from torchvision.transforms import ToTensor
+from neural_processes.modules import BatchNormSequence
+from neural_processes.data.smart_meter import get_smartmeter_df
+
+from neural_processes.utils import ObjectDict
+from neural_processes.lightning import PL_Seq2Seq
+
+
+class Seq2SeqNet(nn.Module):
+    def __init__(self, hparams, _min_std = 0.05):
+        super().__init__()
+        self.hparams = hparams
+        self._min_std = _min_std
+
+
+        self.norm_input = BatchNormSequence(self.hparams.input_size)
+        self.encoder = nn.LSTM(
+            input_size=self.hparams.input_size,
+            hidden_size=self.hparams.hidden_size,
+            batch_first=True,
+            num_layers=self.hparams.lstm_layers,
+            bidirectional=self.hparams.bidirectional,
+            dropout=self.hparams.lstm_dropout,
+        )
+        self.multihead_attn = nn.MultiheadAttention(self.hparams.hidden_size, num_heads=8)
+
+        self.norm_target = BatchNormSequence(self.hparams.input_size_decoder)
+        self.decoder = nn.LSTM(
+            input_size=self.hparams.input_size_decoder,
+            hidden_size=self.hparams.hidden_size,
+            batch_first=True,
+            num_layers=self.hparams.lstm_layers,
+            bidirectional=self.hparams.bidirectional,
+            dropout=self.hparams.lstm_dropout,
+        )
+        self.hidden_out_size = (
+            self.hparams.hidden_size
+            * (self.hparams.bidirectional + 1)
+        )
+        self.mean = nn.Linear(self.hidden_out_size, self.hparams.output_size)
+        self.std = nn.Linear(self.hidden_out_size, self.hparams.output_size)
+        self._use_lvar = False
+
+    def forward(self, context_x, context_y, target_x, target_y=None):
+        x = torch.cat([context_x, context_y], -1)
+
+        # Sometimes input normalisation can be important, an initial batch norm is a nice way to ensure this
+        x = self.norm_input(x)
+        target_x = self.norm_target(target_x)
+
+        _, (h_out, cell) = self.encoder(x)
+        # hidden = [batch size, n layers * n directions, hid dim]
+        # cell = [batch size, n layers * n directions, hid dim]
+
+        # context_x, d_encoded, target_x = k, v, q
+
+        # query, key, value = target_x, context_x, d_encoded
+        attn_output, _ = self.multihead_attn(h_out.permute(1, 0, 2), h_out.permute(1, 0, 2), h_out.permute(1, 0, 2))
+        h_out = attn_output.permute(1, 0, 2).contiguous()
+        attn_output, _ = self.multihead_attn(cell.permute(1, 0, 2), cell.permute(1, 0, 2), cell.permute(1, 0, 2))
+        cell = attn_output.permute(1, 0, 2).contiguous()
+
+        outputs, (_, _) = self.decoder(target_x, (h_out, cell))
+        # output = [batch size, seq len, hid dim * n directions]
+        
+        # outputs: [B, T, num_direction * H]
+        mean = self.mean(outputs)
+        log_sigma = self.std(outputs)
+        if self._use_lvar:
+            log_sigma = torch.clamp(log_sigma, math.log(self._min_std), -math.log(self._min_std))
+            sigma = torch.exp(log_sigma)
+        else:
+            sigma = self._min_std + (1 - self._min_std) * F.softplus(log_sigma)
+        y_dist=torch.distributions.Normal(mean, sigma)
+        
+        # Loss
+        loss_mse = loss_p = None
+        if target_y is not None:
+            loss_mse = F.mse_loss(mean, target_y, reduction='none')
+            if self._use_lvar:
+                loss_p = -log_prob_sigma(target_y, mean, log_sigma)
+            else:
+                loss_p = -y_dist.log_prob(target_y).mean(-1)
+            
+            if self.hparams["context_in_target"]:
+                loss_p[:context_x.size(1)] /= 100
+                loss_mse[:context_x.size(1)] /= 100
+            # # Don't catch loss on context window
+            # mean = mean[:, self.hparams.num_context:]
+            # log_sigma = log_sigma[:, self.hparams.num_context:]
+
+        y_pred = y_dist.rsample if self.training else y_dist.loc
+        return y_pred, dict(loss_p=loss_p.mean(), loss_mse=loss_mse.mean()), dict(log_sigma=log_sigma, dist=y_dist)
+
+
+class LSTMSeq2Seq_PL(PL_Seq2Seq):
+    def __init__(self, hparams,
+        MODEL_CLS=Seq2SeqNet, **kwargs):
+        super().__init__(hparams,
+        MODEL_CLS=MODEL_CLS, **kwargs)
+    
+    @staticmethod
+    def add_suggest(trial):
+        # TODO make label name configurable
+        # TODO make data source configurable
+        trial.suggest_loguniform("learning_rate", 1e-5, 1e-2)
+        trial.suggest_uniform("lstm_dropout", 0, 0.75)
+        trial.suggest_categorical("hidden_size", [1, 2, 4, 8, 16, 32, 64, 128, 256, 512])    
+        trial.suggest_categorical("lstm_layers", [1, 2, 4, 8])    
+        trial.suggest_categorical("bidirectional", [False, True])    
+        
+
+        trial._user_attrs = {
+                'batch_size': 16,
+                'grad_clip': 40,
+                'max_nb_epochs': 200,
+                'num_workers': 4,
+                'num_extra_target': 24*4,
+                'vis_i': '670',
+                'num_context': 24*4,
+                'input_size': 18,
+                'input_size_decoder': 17,
+                'context_in_target': True,
+                'output_size': 1
+        }
+        return trial

neural_processes/models/lstm_std.py

@@ -19,9 +19,9 @@ import io
 import PIL
 from torchvision.transforms import ToTensor
 
-from src.data.smart_meter import get_smartmeter_df
+from neural_processes.data.smart_meter import get_smartmeter_df
 
-from src.utils import ObjectDict
+from neural_processes.utils import ObjectDict
 from torch.utils.data._utils.collate import default_collate
 
 def collate_fn(batch, sample=None):
@@ -112,7 +112,7 @@ class LSTMNet(nn.Module):
         return mean, log_sigma
 
 
-class LSTM_PL(pl.LightningModule):
+class LSTM_PL_STD(pl.LightningModule):
     def __init__(self, hparams):
         # TODO make label name configurable
         # TODO make data source configurable
@@ -126,6 +126,8 @@ class LSTM_PL(pl.LightningModule):
         self._use_lvar = False
         self._min_std = 0.005
 
+        self.default_args = {'bidirectional': False, 'hidden_size_power': 4, 'learning_rate': 0.0010825329363784934, 'lstm_dropout': 0.3905792111699782, 'lstm_layers': 4}
+
     def forward(self, x):
         return self._model(x)
 

neural_processes/models/neural_process/__init__.py

@@ -0,0 +1,2 @@
+from .lightning import PL_NeuralProcess
+from .model import NeuralProcess

neural_processes/models/neural_process/lightning.py

@@ -0,0 +1,82 @@
+import pytorch_lightning as pl
+import torch
+import torch.nn.functional as F
+from argparse import ArgumentParser
+from test_tube import Experiment, HyperOptArgumentParser
+from .model import NeuralProcess
+from neural_processes.lightning import PL_Seq2Seq
+
+
+
+class PL_NeuralProcess(PL_Seq2Seq):
+    def __init__(self, hparams,
+        MODEL_CLS=NeuralProcess, **kwargs):
+        super().__init__(hparams,
+        MODEL_CLS=MODEL_CLS, **kwargs)
+
+    DEFAULT_ARGS = {
+        'attention_dropout': 0,
+        'attention_layers': 2,
+        'batchnorm': False,
+        'det_enc_cross_attn_type': 'multihead',
+        'det_enc_self_attn_type': 'uniform',
+        'dropout': 0,
+        'hidden_dim': 128,
+        'latent_dim': 128,
+        'latent_enc_self_attn_type': 'uniform',
+        'learning_rate': 0.002,
+        'n_decoder_layers': 4,
+        'n_det_encoder_layers': 4,
+        'n_latent_encoder_layers': 2,
+        'num_heads': 8,
+        'use_deterministic_path': True,
+        'use_lvar': True,
+        'use_self_attn': True,
+        'use_rnn': False,
+    }
+
+    @staticmethod
+    def add_suggest(trial):        
+        trial.suggest_loguniform("learning_rate", 1e-5, 1e-2)
+
+        trial.suggest_categorical("hidden_dim", [8*2**i for i in range(8)])
+        trial.suggest_categorical("latent_dim", [8*2**i for i in range(8)])
+        
+        trial.suggest_int("attention_layers", 1, 4)
+        trial.suggest_categorical("n_latent_encoder_layers", [1, 2, 4, 6, 8, 12])
+        trial.suggest_categorical("n_det_encoder_layers", [1, 2, 4, 6, 8, 12])
+        trial.suggest_categorical("n_decoder_layers", [1, 2, 4, 6, 8, 12])
+        trial.suggest_int("num_heads", 8, 8)
+
+        trial.suggest_uniform("dropout", 0, 0.9)
+        trial.suggest_uniform("attention_dropout", 0, 0.9)
+
+        trial.suggest_categorical(
+            "latent_enc_self_attn_type", ['uniform', 'multihead', 'ptmultihead']
+        )
+        trial.suggest_categorical("det_enc_self_attn_type",  ['uniform', 'multihead', 'ptmultihead'])
+        trial.suggest_categorical("det_enc_cross_attn_type", ['uniform', 'multihead', 'ptmultihead'])
+
+        trial.suggest_categorical("batchnorm", [False, True])
+        trial.suggest_categorical("use_self_attn", [False, True])
+        trial.suggest_categorical("use_lvar", [False, True])
+        trial.suggest_categorical("use_deterministic_path", [False, True])
+        trial.suggest_categorical("use_rnn", [True, False])
+
+        trial._user_attrs = {
+            'batch_size': 16,
+            'grad_clip': 40,
+            'max_nb_epochs': 200,
+            'num_workers': 4,
+            'num_context': 24* 4,
+            'vis_i': '670',
+            'num_extra_target': 24*4,
+            'x_dim': 18,
+            'context_in_target': True,
+            'y_dim': 1,
+            'patience': 3,
+            'min_std': 0.005,
+        }        
+        return trial
+
+

neural_processes/models/neural_process/model.py

@@ -4,36 +4,194 @@ import torch.nn.functional as F
 from torch.utils.data import TensorDataset, DataLoader
 import math
 
-from src.models.modules import LatentEncoder, DeterministicEncoder, Decoder
-from src.models.modules import BatchNormSequence
-
-def log_prob_sigma(value, loc, log_scale):
-    """A slightly more stable (not confirmed yet) log prob taking in log_var instead of scale.
-    modified from https://github.com/pytorch/pytorch/blob/2431eac7c011afe42d4c22b8b3f46dedae65e7c0/torch/distributions/normal.py#L65
-    """
-    var = torch.exp(log_scale * 2)
-    return (
-        -((value - loc) ** 2) / (2 * var) - log_scale - math.log(math.sqrt(2 * math.pi))
-    )
+from neural_processes.modules import BatchNormSequence
+from neural_processes.utils import kl_loss_var, log_prob_sigma
 
 
-def kl_loss_var(prior_mu, log_var_prior, post_mu, log_var_post):
-    """
-    Analytical KLD for two gaussians, taking in log_variance instead of scale ( given variance=scale**2) for more stable gradients
-    
-    For version using scale see https://github.com/pytorch/pytorch/blob/master/torch/distributions/kl.py#L398
-    """
+class LatentEncoder(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        hidden_dim=32,
+        latent_dim=32,
+        self_attention_type="dot",
+        n_encoder_layers=3,
+        min_std=0.01,
+        batchnorm=False,
+        dropout=0,
+        attention_dropout=0,
+        use_lvar=False,
+        use_self_attn=False,
+        attention_layers=2,
+        use_lstm=False
+    ):
+        super().__init__()
+        # self._input_layer = nn.Linear(input_dim, hidden_dim)
+        if use_lstm:
+            self._encoder = LSTMBlock(input_dim, hidden_dim, batchnorm=batchnorm, dropout=dropout, num_layers=n_encoder_layers)
+        else:
+            self._encoder = BatchMLP(input_dim, hidden_dim, batchnorm=batchnorm, dropout=dropout, num_layers=n_encoder_layers)
+        if use_self_attn:
+            self._self_attention = Attention(
+                hidden_dim,
+                self_attention_type,
+                attention_layers,
+                rep="identity",
+                dropout=attention_dropout,
+            )
+        self._penultimate_layer = nn.Linear(hidden_dim, hidden_dim)
+        self._mean = nn.Linear(hidden_dim, latent_dim)
+        self._log_var = nn.Linear(hidden_dim, latent_dim)
+        self._min_std = min_std
+        self._use_lvar = use_lvar
+        self._use_lstm = use_lstm
+        self._use_self_attn = use_self_attn
 
-    var_ratio_log = log_var_post - log_var_prior
-    kl_div = (
-        (var_ratio_log.exp() + (post_mu - prior_mu) ** 2) / log_var_prior.exp()
-        - 1.0
-        - var_ratio_log
-    )
-    kl_div = 0.5 * kl_div
-    return kl_div
+    def forward(self, x, y):
+        encoder_input = torch.cat([x, y], dim=-1)
 
-class LatentModel(nn.Module):
+        # Pass final axis through MLP
+        encoded = self._encoder(encoder_input)
+
+        # Aggregator: take the mean over all points
+        if self._use_self_attn:
+            attention_output = self._self_attention(encoded, encoded, encoded)
+            mean_repr = attention_output.mean(dim=1)
+        else:
+            mean_repr = encoded.mean(dim=1)
+
+        # Have further MLP layers that map to the parameters of the Gaussian latent
+        mean_repr = torch.relu(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)
+
+        if self._use_lvar:
+            # Clip it in the log domain, so it can only approach self.min_std, this helps avoid mode collapase
+            # 2 ways, a better but untested way using the more stable log domain, and the way from the deepmind repo
+            log_var = F.logsigmoid(log_var)
+            log_var = torch.clamp(log_var, np.log(self._min_std), -np.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):
+    def __init__(
+        self,
+        input_dim,
+        x_dim,
+        hidden_dim=32,
+        n_d_encoder_layers=3,
+        self_attention_type="dot",
+        cross_attention_type="dot",
+        use_self_attn=False,
+        attention_layers=2,
+        batchnorm=False,
+        dropout=0,
+        attention_dropout=0,
+        use_lstm=False,
+    ):
+        super().__init__()
+        self._use_self_attn = use_self_attn
+        # self._input_layer = nn.Linear(input_dim, hidden_dim)
+        if use_lstm:
+            self._d_encoder = LSTMBlock(input_dim, hidden_dim, batchnorm=batchnorm, dropout=dropout, num_layers=n_d_encoder_layers)
+        else:
+            self._d_encoder = BatchMLP(input_dim, hidden_dim, batchnorm=batchnorm, dropout=dropout, num_layers=n_d_encoder_layers)
+        if use_self_attn:
+            self._self_attention = Attention(
+                hidden_dim,
+                self_attention_type,
+                attention_layers,
+                rep="identity",
+                dropout=attention_dropout,
+            )
+        self._cross_attention = Attention(
+            hidden_dim,
+            cross_attention_type,
+            x_dim=x_dim,
+            attention_layers=attention_layers,
+        )
+
+    def forward(self, context_x, context_y, target_x):
+        # Concatenate x and y along the filter axes
+        d_encoder_input = torch.cat([context_x, context_y], dim=-1)
+
+        # Pass final axis through MLP
+        d_encoded = self._d_encoder(d_encoder_input)
+
+        if self._use_self_attn:
+            d_encoded = self._self_attention(d_encoded, d_encoded, d_encoded)
+
+        # Apply attention as mean aggregation
+        h = self._cross_attention(context_x, d_encoded, target_x)
+
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        x_dim,
+        y_dim,
+        hidden_dim=32,
+        latent_dim=32,
+        n_decoder_layers=3,
+        use_deterministic_path=True,
+        min_std=0.01,
+        use_lvar=False,
+        batchnorm=False,
+        dropout=0,
+        use_lstm=False,
+    ):
+        super(Decoder, self).__init__()
+        self._target_transform = nn.Linear(x_dim, hidden_dim)
+        if use_deterministic_path:
+            hidden_dim_2 = 2 * hidden_dim + latent_dim
+        else:
+            hidden_dim_2 = hidden_dim + latent_dim
+            
+        if use_lstm:
+            self._decoder = LSTMBlock(hidden_dim_2, hidden_dim_2, batchnorm=batchnorm, dropout=dropout, num_layers=n_decoder_layers)
+        else:
+            self._decoder = BatchMLP(hidden_dim_2, hidden_dim_2, batchnorm=batchnorm, dropout=dropout, num_layers=n_decoder_layers)
+        self._mean = nn.Linear(hidden_dim_2, y_dim)
+        self._std = nn.Linear(hidden_dim_2, y_dim)
+        self._use_deterministic_path = use_deterministic_path
+        self._min_std = min_std
+        self._use_lvar = use_lvar
+
+    def forward(self, r, z, target_x):
+        # concatenate target_x and representation
+        x = self._target_transform(target_x)
+
+        if self._use_deterministic_path:
+            z = torch.cat([r, z], dim=-1)
+
+        r = torch.cat([z, x], dim=-1)
+
+        r = self._decoder(r)
+
+        # Get the mean and the variance
+        mean = self._mean(r)
+        log_sigma = self._std(r)
+
+        # Bound or clamp the variance
+        if self._use_lvar:
+            log_sigma = torch.clamp(log_sigma, math.log(self._min_std), -math.log(self._min_std))
+            sigma = torch.exp(log_sigma)
+        else:
+            sigma = self._min_std + (1 - self._min_std) * F.softplus(log_sigma)
+
+        dist = torch.distributions.Normal(mean, sigma)
+        return dist, log_sigma
+
+
+class NeuralProcess(nn.Module):
     def __init__(self,
                  x_dim, # features in input
                  y_dim, # number of features in output
@@ -61,7 +219,7 @@ class LatentModel(nn.Module):
                  **kwargs,
                 ):
 
-        super(LatentModel, self).__init__()
+        super(NeuralProcess, self).__init__()
 
         self._use_rnn = use_rnn
         self.context_in_target = context_in_target

neural_processes/models/transformer.py

@@ -0,0 +1,165 @@
+import os
+import numpy as np
+import pandas as pd
+import torch
+import optuna
+from tqdm.auto import tqdm
+from torch import nn
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+
+from neural_processes.lightning import PL_Seq2Seq
+
+class NetTransformer(nn.Module):
+    def __init__(self, hparams):
+        super().__init__()
+        hparams["nlayers"] = int(2 ** hparams["nlayers_power"])
+        hparams["hidden_size"] = int(2**hparams["hidden_size_power"])
+        hparams["hidden_out_size"] = int(2 ** hparams["hidden_out_size_power"])
+        hparams["nhead"] = int(2 ** hparams["nhead_power"])
+        logger.debug(f"{type(self)} hparams {hparams}")
+        self.hparams = hparams
+
+        hidden_out_size = self.hparams.hidden_out_size
+        enc_input_size = self.hparams.input_size + self.hparams.output_size
+        # self.enc_norm = BatchNormSequence(enc_input_size)
+        self.enc_emb = nn.Linear(enc_input_size, hidden_out_size)
+        encoder_norm = nn.LayerNorm(hidden_out_size)
+        layer_enc = nn.TransformerEncoderLayer(
+            d_model=hidden_out_size,
+            dim_feedforward=self.hparams.hidden_size,
+            dropout=self.hparams.attention_dropout,
+            nhead=self.hparams.nhead,
+            # activation
+        )
+        self.encoder = nn.TransformerEncoder(
+            layer_enc,
+            num_layers=self.hparams.nlayers,
+            norm=encoder_norm
+        )
+
+        # self.dec_norm = BatchNormSequence(self.hparams.input_size)
+        # self.dec_emb = nn.Linear(self.hparams.input_size, hidden_out_size)
+        # layer_dec = nn.TransformerDecoderLayer(
+        #     d_model=hidden_out_size,
+        #     dim_feedforward=self.hparams.hidden_size,
+        #     dropout=self.hparams.attention_dropout,
+        #     nhead=self.hparams.nhead,
+        # )
+        # decoder_norm = nn.LayerNorm(hidden_out_size)
+        # self.decoder = nn.TransformerDecoder(
+        #     layer_dec,
+        #     num_layers=self.hparams.nlayers,
+        #     norm=decoder_norm
+        # )
+        self.mean = nn.Linear(hidden_out_size, self.hparams.output_size)
+
+    def forward(self, context_x, context_y, target_x, target_y=None):
+        device = next(self.parameters()).device
+        target_y_fake = torch.ones(context_y.shape).float().to(device) * self.hparams.nan_value
+        context = torch.cat([context_x, context_y], -1).detach()
+        target = torch.cat([target_x, target_y_fake], -1).detach()
+        x = torch.cat([context, target * 1], 1).detach()
+
+        # Masks
+        x_mask = torch.isfinite(x) & (x!=self.hparams.nan_value)
+        x[~x_mask] = 0
+        x = x.detach()
+        x_key_padding_mask = ~x_mask.any(-1)
+        # print('x_key_padding_mask', x_mask.float().mean())
+        # print(x.shape, 'x1')
+        x = self.enc_emb(x).permute(1, 0, 2)
+        # print(x.shape, 'x2')
+        # Size([C, B, emb_dim])
+        outputs = self.encoder(x, src_key_padding_mask=x_key_padding_mask).permute(1, 0, 2)
+        # print(outputs.shape, 'outputs')
+
+        # Seems to help a little, especially with extrapolating out of bounds
+        steps = target_y.shape[1]
+        mean = softnorm(self.mean(outputs))
+        mean_target = mean[:, -steps:, :]
+        mean_context = mean[:, :-steps, :]
+
+        loss = None
+        if target_y is not None:
+            y = torch.cat([context_y, target_y], 1)
+            y_mask = torch.isfinite(y) & (y!=self.hparams.nan_value)
+            y[~y_mask] = 0
+            y = y.detach()
+
+            
+            loss_scale = 100
+            # loss = F.mse_loss(mean * loss_scale, y * loss_scale, reduction='none') / loss_scale
+
+            loss_target = F.mse_loss(mean_target * loss_scale, y[:, -steps:, :] * loss_scale, reduction='none') / loss_scale
+            loss_context = F.mse_loss(mean_context * loss_scale, y[:, :-steps, :] * loss_scale, reduction='none') / loss_scale
+
+            y_mask_target = y_mask[:, -steps:, :].detach()
+            y_mask_context = y_mask[:, :-steps, :].detach()
+            # loss_target = loss[:, -steps:, :]
+            # loss_context = loss[:, :-steps, :]
+            # print(0, loss_context.sum(), loss_target.sum())
+            
+            weight = (torch.arange(loss_target.shape[1]) + 0.5).float().to(device)[None,:, None]
+            # weight /= weight.sum()
+            # print(1.0, loss_context.sum(), loss_target.sum())
+            loss_target = loss_target / torch.sqrt(weight)  # We want to weight nearer stuff more
+            # print(1.5, loss_context.sum(), y_mask_context.sum(), loss_target.sum(), y_mask_target.sum(), (loss_context * y_mask_context).sum())
+            loss_context = (loss_context * y_mask_context.float()).sum() / (y_mask_context.sum() + 1.)
+            loss_target = (loss_target * y_mask_target.float()).sum() / (y_mask_target.sum()+1.)  # Mean over unmasked ones
+            # print(2, loss_context.sum(), loss_target.sum())
+            
+            # Perhaps predicting the past, as a secondary loss will help
+            loss =  loss_context /100. + loss_target
+
+            assert torch.isfinite(loss)
+
+        return mean_target, dict(loss=loss), dict()
+
+class PL_Transformer(PL_Seq2Seq):
+    def __init__(self, hparams,
+        MODEL_CLS=NetTransformer, **kwargs):
+        super().__init__(hparams,
+        MODEL_CLS=MODEL_CLS, **kwargs)
+
+        self.default_args = {'attention_dropout': 0.4151003234623061, 'hidden_out_size_power': 2.0, 'hidden_size_power': 2.0, 'learning_rate': 0.0026738884132767185, 'nhead_power': 1.0, 'nlayers_power': 1.0}
+
+    @staticmethod
+    def add_suggest(trial: optuna.Trial, user_attrs={}):
+        """
+        Add hyperparam ranges to an optuna trial and typical user attrs.
+        
+        Usage:
+            trial = optuna.trial.FixedTrial(
+                params={         
+                    'hidden_size': 128,
+                }
+            )
+            trial = add_suggest(trial)
+            trainer = pl.Trainer()
+            model = LSTM_PL(dict(**trial.params, **trial.user_attrs), dataset_train,
+                            dataset_test, cache_base_path, norm)
+            trainer.fit(model)
+        """
+        trial.suggest_loguniform("learning_rate", 1e-5, 1e-2)
+        trial.suggest_uniform("attention_dropout", 0, 0.9)
+        trial.suggest_discrete_uniform(
+            "hidden_size_power", 2, 10, 1
+        )
+        trial.suggest_discrete_uniform("hidden_out_size_power", 2, 9, 1)
+        trial.suggest_discrete_uniform("nhead_power", 1, 4, 1)
+        trial.suggest_discrete_uniform("nlayers_power", 1, 5, 1)      
+
+        user_attrs_default = {
+            "batch_size": 16,
+            "grad_clip": 40,
+            "max_nb_epochs": 200,
+            "num_workers": 4,
+            "vis_i": 670,
+            "input_size": 6,
+            "output_size": 1,
+            "label_steps": 24,
+        }
+        [trial.set_user_attr(k, v) for k, v in user_attrs_default.items()]
+        [trial.set_user_attr(k, v) for k, v in user_attrs.items()]
+        return trial

neural_processes/models/transformer_seq2seq.py

@@ -8,9 +8,9 @@ from torch.nn import functional as F
 from torch.utils.data import DataLoader
 from torchvision.datasets import MNIST
 from test_tube import Experiment, HyperOptArgumentParser
-from src.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df
+from neural_processes.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df
 import torchvision.transforms as transforms
-from src.plot import plot_from_loader_to_tensor, plot_from_loader
+from neural_processes.plot import plot_from_loader_to_tensor, plot_from_loader
 from argparse import ArgumentParser
 import json
 import pytorch_lightning as pl
@@ -22,10 +22,11 @@ import PIL
 import optuna
 from torchvision.transforms import ToTensor
 
-from src.data.smart_meter import get_smartmeter_df
-from src.models.modules import BatchNormSequence
+from neural_processes.data.smart_meter import get_smartmeter_df
+from neural_processes.modules import BatchNormSequence
 
-from src.utils import ObjectDict
+from neural_processes.utils import ObjectDict
+from neural_processes.lightning import PL_Seq2Seq
 
 def log_prob_sigma(value, loc, log_scale):
     """A slightly more stable (not confirmed yet) log prob taking in log_var instead of scale.
@@ -136,138 +137,13 @@ class TransformerSeq2SeqNet(nn.Module):
         return y_pred, dict(loss_p=loss_p.mean(), loss_mse=loss_mse.mean()), dict(log_sigma=log_sigma, dist=y_dist)
 
 
-class TransformerSeq2Seq_PL(pl.LightningModule):
-    def __init__(self, hparams):
-        # TODO make label name configurable
-        # TODO make data source configurable
-        super().__init__()
-        self.hparams = ObjectDict()
-        self.hparams.update(
-            hparams.__dict__ if hasattr(hparams, "__dict__") else hparams
-        )
-        self.model = TransformerSeq2SeqNet(self.hparams)
-        self._dfs = None
+class TransformerSeq2Seq_PL(PL_Seq2Seq):
 
-    def forward(self, context_x, context_y, target_x, target_y):
-        return self.model(context_x, context_y, target_x, target_y)
-
-    def training_step(self, batch, batch_idx):
-        # REQUIRED
-        assert all(torch.isfinite(d).all() for d in batch)
-        context_x, context_y, target_x, target_y = batch
-        y_dist, losses, extra = self.forward(context_x, context_y, target_x, target_y)
-        loss = losses['loss_p'] # + loss_mse
-        tensorboard_logs = {
-            "train/loss": loss,
-            'train/loss_mse': losses['loss_mse'],
-            "train/loss_p": losses['loss_p'],
-            "train/sigma": torch.exp(extra['log_sigma']).mean()}
-        return {"loss": loss, "log": tensorboard_logs}
-
-    def validation_step(self, batch, batch_idx):
-        context_x, context_y, target_x, target_y = batch
-        assert all(torch.isfinite(d).all() for d in batch)
-        y_dist, losses, extra = self.forward(context_x, context_y, target_x, target_y)
-        loss = losses['loss_p'] # + loss_mse
-        tensorboard_logs = {
-            "val_loss": loss,
-            'val/loss_mse': losses['loss_mse'],
-            "val/loss_p": losses['loss_p'],
-            "val/sigma": torch.exp(extra['log_sigma']).mean()}
-        return {"val_loss": loss, "log": tensorboard_logs}
-
-    def validation_end(self, outputs):
-        if int(self.hparams["vis_i"]) > 0:
-            self.show_image()
-
-        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
-        keys = outputs[0]["log"].keys()
-        tensorboard_logs = {
-            k: torch.stack([x["log"][k] for x in outputs if k in x["log"]]).mean()
-            for k in keys
-        }
-        tensorboard_logs_str = {k: f"{v}" for k, v in tensorboard_logs.items()}
-        print(f"step {self.trainer.global_step}, {tensorboard_logs_str}")
-        assert torch.isfinite(avg_loss)
-        return {"avg_val_loss": avg_loss, "log": tensorboard_logs}
-
-
-    def show_image(self):        
-        # https://github.com/PytorchLightning/pytorch-lightning/blob/f8d9f8f/pytorch_lightning/core/lightning.py#L293
-        loader = self.val_dataloader()
-        vis_i = min(int(self.hparams["vis_i"]), len(loader.dataset))
-        # print('vis_i', vis_i)
-        if isinstance(self.hparams["vis_i"], str):
-            image = plot_from_loader(loader, self, i=int(vis_i))
-            plt.show()
-        else:
-            image = plot_from_loader_to_tensor(loader, self, i=vis_i)
-            self.logger.experiment.add_image('val/image', image, self.trainer.global_step)
-
-    def test_step(self, *args, **kwargs):
-        return self.validation_step(*args, **kwargs)
-
-    def test_end(self, *args, **kwargs):
-        return self.validation_end(*args, **kwargs)
-
-    def configure_optimizers(self):
-        optim = torch.optim.Adam(self.parameters(), lr=self.hparams["learning_rate"])
-        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
-            optim, patience=self.hparams["patience"], verbose=True, min_lr=1e-7
-        )  # note early stopping has patience 3
-        return [optim], [scheduler]
-
-    def _get_cache_dfs(self):
-        if self._dfs is None:
-            df_train, df_val, df_test = get_smartmeter_df()
-            self._dfs = dict(df_train=df_train, df_val=df_val, df_test=df_test)
-        return self._dfs
-
-    @pl.data_loader
-    def train_dataloader(self):
-        df_train = self._get_cache_dfs()['df_train']
-        data_train = SmartMeterDataSet(
-            df_train, self.hparams["num_context"], self.hparams["num_extra_target"]
-        )
-        return torch.utils.data.DataLoader(
-            data_train,
-            batch_size=self.hparams["batch_size"],
-            shuffle=True,
-            collate_fn=collate_fns(
-                self.hparams["num_context"], self.hparams["num_extra_target"], sample=True, context_in_target=self.hparams["context_in_target"]
-            ),
-            num_workers=self.hparams["num_workers"],
-        )
-
-    @pl.data_loader
-    def val_dataloader(self):
-        df_test = self._get_cache_dfs()['df_val']
-        data_test = SmartMeterDataSet(
-            df_test, self.hparams["num_context"], self.hparams["num_extra_target"]
-        )
-        return torch.utils.data.DataLoader(
-            data_test,
-            batch_size=self.hparams["batch_size"],
-            shuffle=False,
-            collate_fn=collate_fns(
-                self.hparams["num_context"], self.hparams["num_extra_target"], sample=False, context_in_target=self.hparams["context_in_target"]
-            ),
-        )
-
-    @pl.data_loader
-    def test_dataloader(self):
-        df_test = self._get_cache_dfs()['df_test']
-        data_test = SmartMeterDataSet(
-            df_test, self.hparams["num_context"], self.hparams["num_extra_target"]
-        )
-        return torch.utils.data.DataLoader(
-            data_test,
-            batch_size=self.hparams["batch_size"],
-            shuffle=False,
-            collate_fn=collate_fns(
-                self.hparams["num_context"], self.hparams["num_extra_target"], sample=False, context_in_target=self.hparams["context_in_target"]
-            ),
-        )
+    def __init__(self, hparams,
+        MODEL_CLS=TransformerSeq2SeqNet, **kwargs):
+        super().__init__(hparams,
+        MODEL_CLS=MODEL_CLS, **kwargs)
+        self.default_args = {'agg': 'mean', 'attention_dropout': 0.12013231612195126, 'hidden_out_size_power': 4.0, 'hidden_size_power': 7.0, 'learning_rate': 0.0022924639229335475, 'nhead_power': 2.0, 'nlayers_power': 4.0}
 
     @staticmethod
     def add_suggest(trial: optuna.Trial):

neural_processes/modules/__init__.py

@@ -0,0 +1,2 @@
+from .modules import BatchMLP, BatchNormSequence, LSTMBlock
+from .attention import Attention, AttnLinear

neural_processes/modules/attention.py

@@ -0,0 +1,124 @@
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+class AttnLinear(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.linear = nn.Linear(in_channels, out_channels, bias=False)
+        torch.nn.init.normal_(self.linear.weight, std=in_channels ** -0.5)
+
+    def forward(self, x):
+        x = self.linear(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        hidden_dim,
+        attention_type,
+        attention_layers=2,
+        n_heads=8,
+        x_dim=1,
+        rep="mlp",
+        dropout=0,
+        batchnorm=False,
+    ):
+        super().__init__()
+        self._rep = rep
+
+        if self._rep == "mlp":
+            self.batch_mlp_k = BatchMLP(
+                x_dim,
+                hidden_dim,
+                attention_layers,
+                dropout=dropout,
+                batchnorm=batchnorm,
+            )
+            self.batch_mlp_q = BatchMLP(
+                x_dim,
+                hidden_dim,
+                attention_layers,
+                dropout=dropout,
+                batchnorm=batchnorm,
+            )
+
+        if attention_type == "uniform":
+            self._attention_func = self._uniform_attention
+        elif attention_type == "laplace":
+            self._attention_func = self._laplace_attention
+        elif attention_type == "dot":
+            self._attention_func = self._dot_attention
+        elif attention_type == "multihead":
+            self._W_k = nn.ModuleList(
+                [AttnLinear(hidden_dim, hidden_dim) for _ in range(n_heads)]
+            )
+            self._W_v = nn.ModuleList(
+                [AttnLinear(hidden_dim, hidden_dim) for _ in range(n_heads)]
+            )
+            self._W_q = nn.ModuleList(
+                [AttnLinear(hidden_dim, hidden_dim) for _ in range(n_heads)]
+            )
+            self._W = AttnLinear(n_heads * hidden_dim, hidden_dim)
+            self._attention_func = self._multihead_attention
+            self.n_heads = n_heads
+        elif attention_type == "ptmultihead":
+            self._W = torch.nn.MultiheadAttention(
+                hidden_dim, n_heads, bias=False, dropout=dropout
+            )
+            self._attention_func = self._pytorch_multihead_attention
+        else:
+            raise NotImplementedError
+
+    def forward(self, k, v, q):
+        if self._rep == "mlp":
+            k = self.batch_mlp_k(k)
+            q = self.batch_mlp_q(q)
+        rep = self._attention_func(k, v, q)
+        return rep
+
+    def _uniform_attention(self, k, v, q):
+        total_points = q.shape[1]
+        rep = torch.mean(v, dim=1, keepdim=True)
+        rep = rep.repeat(1, total_points, 1)
+        return rep
+
+    def _laplace_attention(self, k, v, q, scale=0.5):
+        k_ = k.unsqueeze(1)
+        v_ = v.unsqueeze(2)
+        unnorm_weights = torch.abs((k_ - v_) * scale)
+        unnorm_weights = unnorm_weights.sum(dim=-1)
+        weights = torch.softmax(unnorm_weights, dim=-1)
+        rep = torch.einsum("bik,bkj->bij", weights, v)
+        return rep
+
+    def _dot_attention(self, k, v, q):
+        scale = q.shape[-1] ** 0.5
+        unnorm_weights = torch.einsum("bjk,bik->bij", k, q) / scale
+        weights = torch.softmax(unnorm_weights, dim=-1)
+
+        rep = torch.einsum("bik,bkj->bij", weights, v)
+        return rep
+
+    def _multihead_attention(self, k, v, q):
+        outs = []
+        for i in range(self.n_heads):
+            k_ = self._W_k[i](k)
+            v_ = self._W_v[i](v)
+            q_ = self._W_q[i](q)
+            out = self._dot_attention(k_, v_, q_)
+            outs.append(out)
+        outs = torch.stack(outs, dim=-1)
+        outs = outs.view(outs.shape[0], outs.shape[1], -1)
+        rep = self._W(outs)
+        return rep
+
+    def _pytorch_multihead_attention(self, k, v, q):
+        # Pytorch multiheaded attention takes inputs if diff order and permutation
+        q = q.permute(1, 0, 2)
+        k = k.permute(1, 0, 2)
+        v = v.permute(1, 0, 2)
+        o = self._W(q, k, v)[0]
+        return o.permute(1, 0, 2)

neural_processes/modules/modules.py

@@ -0,0 +1,107 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+import numpy as np
+# from .attention import Attention as PtAttention
+
+
+class LSTMBlock(nn.Module):
+    def __init__(
+        self, in_channels, out_channels, dropout=0, batchnorm=False, bias=False, num_layers=1
+    ):
+        super().__init__()
+        self._lstm = nn.LSTM(
+                input_size=in_channels,
+                hidden_size=out_channels,
+                num_layers=num_layers,
+                dropout=dropout,
+                batch_first=True,
+                bias=bias
+        )
+
+    def forward(self, x):
+        return self._lstm(x)[0]
+
+
+class BatchNormSequence(nn.Module):
+    """Applies batch norm on features of a batch first sequence."""
+    def __init__(
+        self, out_channels
+    ):
+        super().__init__()
+        self.norm = nn.BatchNorm1d(out_channels)
+
+    def forward(self, x):
+        # x.shape is (Batch, Sequence, Channels)
+        # Now we want to apply batchnorm and dropout to the channels. So we put it in shape
+        # (Batch, Channels, Sequence) so we can use BatchNorm1d
+        x = x.permute(0, 2, 1)
+        x = self.norm(x)
+        return x.permute(0, 2, 1)
+
+class NPBlockRelu2d(nn.Module):
+    """Block for Neural Processes."""
+
+    def __init__(
+        self, in_channels, out_channels, dropout=0, batchnorm=False, bias=False
+    ):
+        super().__init__()
+        self.linear = nn.Linear(in_channels, out_channels, bias=bias)
+        self.act = nn.ReLU()
+        self.dropout = nn.Dropout2d(dropout)
+        self.norm = nn.BatchNorm2d(out_channels) if batchnorm else False
+
+    def forward(self, x):
+        # x.shape is (Batch, Sequence, Channels)
+        # We pass a linear over it which operates on the Channels
+        x = self.act(self.linear(x))
+
+        # Now we want to apply batchnorm and dropout to the channels. So we put it in shape
+        # (Batch, Channels, Sequence, None) so we can use Dropout2d & BatchNorm2d
+        x = x.permute(0, 2, 1)[:, :, :, None]
+
+        if self.norm:
+            x = self.norm(x)
+
+        x = self.dropout(x)
+        return x[:, :, :, 0].permute(0, 2, 1)
+
+
+class BatchMLP(nn.Module):
+    """Apply MLP to the final axis of a 3D tensor (reusing already defined MLPs).
+
+    Args:
+        input: input tensor of shape [B,n,d_in].
+        output_sizes: An iterable containing the output sizes of the MLP as defined 
+            in `basic.Linear`.
+    Returns:
+        tensor of shape [B,n,d_out] where d_out=output_size
+    """
+
+    def __init__(
+        self, input_size, output_size, num_layers=2, dropout=0, batchnorm=False
+    ):
+        super().__init__()
+        self.input_size = input_size
+        self.output_size = output_size
+        self.num_layers = num_layers
+
+        self.initial = NPBlockRelu2d(
+            input_size, output_size, dropout=dropout, batchnorm=batchnorm
+        )
+        self.encoder = nn.Sequential(
+            *[
+                NPBlockRelu2d(
+                    output_size, output_size, dropout=dropout, batchnorm=batchnorm
+                )
+                for _ in range(num_layers - 2)
+            ]
+        )
+        self.final = nn.Linear(output_size, output_size)
+
+    def forward(self, x):
+        x = self.initial(x)
+        x = self.encoder(x)
+        return self.final(x)
+

neural_processes/train.py

@@ -8,6 +8,7 @@ import torch
 from .dict_logger import DictLogger
 from .utils import PyTorchLightningPruningCallback
 from .plot import plot_from_loader
+from .logger import logger
 
 
 def main(
@@ -89,10 +90,24 @@ def run_trial(
 ):
     print(f"now run `tensorboard --logdir {MODEL_DIR}`")
     (MODEL_DIR / name).mkdir(parents=True, exist_ok=True)
+
+    if getattr(PL_MODEL_CLS, 'default_args', None):
+        # add default args
+        params = {**PL_MODEL_CLS.default_args, **params}
+    else:
+        logger.warning(f"No default args on {PL_MODEL_CLS}")
+
+    # Make trial
     trial = optuna.trial.FixedTrial(params=params)
     trial = PL_MODEL_CLS.add_suggest(trial)
+
+    # Add auto number
     trial = add_number(trial, MODEL_DIR / name)
+
+    # Add user attributes
     trial._user_attrs.update(user_attrs)
+    print(trial)
+
     model, trainer = main(
         trial, PL_MODEL_CLS, name=name, MODEL_DIR=MODEL_DIR, train=False, prune=False
     )

neural_processes/utils.py

@@ -77,3 +77,29 @@ class ObjectDict(dict):
     def __dict__(self):
         return dict(self)
 
+
+def log_prob_sigma(value, loc, log_scale):
+    """A slightly more stable (not confirmed yet) log prob taking in log_var instead of scale.
+    modified from https://github.com/pytorch/pytorch/blob/2431eac7c011afe42d4c22b8b3f46dedae65e7c0/torch/distributions/normal.py#L65
+    """
+    var = torch.exp(log_scale * 2)
+    return (
+        -((value - loc) ** 2) / (2 * var) - log_scale - math.log(math.sqrt(2 * math.pi))
+    )
+
+
+def kl_loss_var(prior_mu, log_var_prior, post_mu, log_var_post):
+    """
+    Analytical KLD for two gaussians, taking in log_variance instead of scale ( given variance=scale**2) for more stable gradients
+    
+    For version using scale see https://github.com/pytorch/pytorch/blob/master/torch/distributions/kl.py#L398
+    """
+
+    var_ratio_log = log_var_post - log_var_prior
+    kl_div = (
+        (var_ratio_log.exp() + (post_mu - prior_mu) ** 2) / log_var_prior.exp()
+        - 1.0
+        - var_ratio_log
+    )
+    kl_div = 0.5 * kl_div
+    return kl_div

smartmeters-ANP-RNN-Copy1.ipynb

@@ -110,11 +110,11 @@
    },
    "outputs": [],
    "source": [
-    "from src.models.model import LatentModel\n",
-    "from src.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df\n",
-    "from src.plot import plot_from_loader\n",
-    "from src.models.lightning_anp import LatentModelPL\n",
-    "from src.dict_logger import DictLogger"
+    "from neural_processes.models.model import LatentModel\n",
+    "from neural_processes.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df\n",
+    "from neural_processes.plot import plot_from_loader\n",
+    "from neural_processes.models.lightning_anp import LatentModelPL\n",
+    "from neural_processes.dict_logger import DictLogger"
    ]
   },
   {

smartmeters-ANP-RNN-mcdropout.ipynb

@@ -105,13 +105,13 @@
    },
    "outputs": [],
    "source": [
-    "from src.models.model import LatentModel\n",
-    "from src.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df\n",
-    "from src.plot import plot_from_loader\n",
-    "from src.models.lightning_anp import LatentModelPL\n",
-    "from src.dict_logger import DictLogger\n",
-    "from src.train import main, objective, add_number, run_trial\n",
-    "from src.utils import init_random_seed"
+    "from neural_processes.models.model import LatentModel\n",
+    "from neural_processes.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df\n",
+    "from neural_processes.plot import plot_from_loader\n",
+    "from neural_processes.models.lightning_anp import LatentModelPL\n",
+    "from neural_processes.dict_logger import DictLogger\n",
+    "from neural_processes.train import main, objective, add_number, run_trial\n",
+    "from neural_processes.utils import init_random_seed"
    ]
   },
   {
@@ -289,7 +289,7 @@
    },
    "outputs": [],
    "source": [
-    "from src.plot import plot_rows\n",
+    "from neural_processes.plot import plot_rows\n",
     "\n",
     "\n",
     "def eval_mc(model, loader, i):\n",
@@ -392,7 +392,7 @@
     " 'use_self_attn': True, \n",
     " 'use_rnn': False, \n",
     "}\n",
-    "default_attrs = {\n",
+    "default_user_attrs = {\n",
     "    'context_in_target': True,\n",
     "    'x_dim': 17,\n",
     "     'y_dim': 1,\n",
@@ -2835,7 +2835,7 @@
     "         'use_rnn': True,\n",
     "         'vis_i': 670\n",
     "        },\n",
-    "        user_attrs = default_attrs,\n",
+    "        user_attrs = default_user_attrs,\n",
     "        PL_MODEL_CLS=LatentModelPL\n",
     "    )\n",
     "    \n",
@@ -3256,7 +3256,7 @@
    },
    "outputs": [],
    "source": [
-    "from src.plot import plot_rows\n",
+    "from neural_processes.plot import plot_rows\n",
     "loader = model.val_dataloader()\n",
     "device = next(model.parameters()).device\n",
     "\n",

smartmeters-lstm-deterministic.ipynb

@@ -85,13 +85,13 @@
    },
    "outputs": [],
    "source": [
-    "from src.models.model import LatentModel\n",
-    "from src.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df\n",
-    "# from src.plot import plot_from_loader\n",
-    "from src.models.lstm import SequenceDfDataSet, LSTM_PL, plot_from_loader\n",
-    "from src.dict_logger import DictLogger\n",
-    "from src.utils import PyTorchLightningPruningCallback\n",
-    "from src.train import main, objective, add_number, run_trial"
+    "from neural_processes.models.model import LatentModel\n",
+    "from neural_processes.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df\n",
+    "# from neural_processes.plot import plot_from_loader\n",
+    "from neural_processes.models.lstm import SequenceDfDataSet, LSTM_PL, plot_from_loader\n",
+    "from neural_processes.dict_logger import DictLogger\n",
+    "from neural_processes.utils import PyTorchLightningPruningCallback\n",
+    "from neural_processes.train import main, objective, add_number, run_trial"
    ]
   },
   {

smartmeters-lstm-seq2seq.ipynb

@@ -96,13 +96,13 @@
     }
    ],
    "source": [
-    "from src.models.model import LatentModel\n",
-    "from src.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df\n",
-    "from src.plot import plot_from_loader\n",
-    "from src.models.lstm_seqseq import LSTMSeq2Seq_PL\n",
-    "from src.dict_logger import DictLogger\n",
-    "from src.utils import PyTorchLightningPruningCallback\n",
-    "from src.train import main, objective, add_number, run_trial"
+    "from neural_processes.models.model import LatentModel\n",
+    "from neural_processes.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df\n",
+    "from neural_processes.plot import plot_from_loader\n",
+    "from neural_processes.models.lstm_seqseq import LSTMSeq2Seq_PL\n",
+    "from neural_processes.dict_logger import DictLogger\n",
+    "from neural_processes.utils import PyTorchLightningPruningCallback\n",
+    "from neural_processes.train import main, objective, add_number, run_trial"
    ]
   },
   {

smartmeters-lstm_std.ipynb

@@ -96,13 +96,13 @@
     }
    ],
    "source": [
-    "from src.models.model import LatentModel\n",
-    "from src.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df\n",
-    "# from src.plot import plot_from_loader\n",
-    "from src.models.lstm_std import SequenceDfDataSet, LSTM_PL, plot_from_loader\n",
-    "from src.dict_logger import DictLogger\n",
-    "from src.utils import PyTorchLightningPruningCallback\n",
-    "from src.train import main, objective, add_number, run_trial"
+    "from neural_processes.models.model import LatentModel\n",
+    "from neural_processes.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df\n",
+    "# from neural_processes.plot import plot_from_loader\n",
+    "from neural_processes.models.lstm_std import SequenceDfDataSet, LSTM_PL, plot_from_loader\n",
+    "from neural_processes.dict_logger import DictLogger\n",
+    "from neural_processes.utils import PyTorchLightningPruningCallback\n",
+    "from neural_processes.train import main, objective, add_number, run_trial"
    ]
   },
   {

smartmeters-transformer-seq2seq.ipynb

@@ -96,13 +96,13 @@
     }
    ],
    "source": [
-    "from src.models.model import LatentModel\n",
-    "from src.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df\n",
-    "from src.plot import plot_from_loader\n",
-    "from src.models.transformer_seq2seq import TransformerSeq2Seq_PL\n",
-    "from src.dict_logger import DictLogger\n",
-    "from src.utils import PyTorchLightningPruningCallback\n",
-    "from src.train import main, objective, add_number, run_trial"
+    "from neural_processes.models.model import LatentModel\n",
+    "from neural_processes.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df\n",
+    "from neural_processes.plot import plot_from_loader\n",
+    "from neural_processes.models.transformer_seq2seq import TransformerSeq2Seq_PL\n",
+    "from neural_processes.dict_logger import DictLogger\n",
+    "from neural_processes.utils import PyTorchLightningPruningCallback\n",
+    "from neural_processes.train import main, objective, add_number, run_trial"
    ]
   },
   {

src/models/lightning_anp.py

@@ -1,217 +0,0 @@
-import pytorch_lightning as pl
-import torch
-import torch.nn.functional as F
-from argparse import ArgumentParser
-from test_tube import Experiment, HyperOptArgumentParser
-from src.models.model import LatentModel
-from src.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df
-from src.plot import plot_from_loader_to_tensor, plot_from_loader
-from src.utils import ObjectDict
-from matplotlib import pyplot as plt
-
-
-
-class LatentModelPL(pl.LightningModule):
-    def __init__(self, hparams):
-        super().__init__()
-        self.hparams = ObjectDict()
-        self.hparams.update(hparams.__dict__ if hasattr(hparams, '__dict__') else hparams)
-        self.model = LatentModel(**self.hparams)
-        self._dfs = None
-        self.train_logs = [] 
-
-    def forward(self, context_x, context_y, target_x, target_y):
-        return self.model(context_x, context_y, target_x, target_y)
-
-    def training_step(self, batch, batch_idx):
-        assert all(torch.isfinite(d).all() for d in batch)
-        context_x, context_y, target_x, target_y = batch
-        y_pred, losses, extra = self.forward(context_x, context_y, target_x, target_y)
-        y_std = extra['dist'].scale
-        loss = losses['loss'].mean()
-
-        tensorboard_logs = {
-            "train_loss": loss,
-            "train/kl": losses['loss_kl'].mean(),
-            "train/std": y_std.mean(),
-            "train/mse": losses['loss_mse'].mean(),
-        }
-        assert torch.isfinite(loss)
-        self.train_logs.append(tensorboard_logs)
-        return {"loss": loss, "log": tensorboard_logs}
-
-    def validation_step(self, batch, batch_idx):
-        assert all(torch.isfinite(d).all() for d in batch)
-        context_x, context_y, target_x, target_y = batch
-        y_pred, losses, extra = self.forward(context_x, context_y, target_x, target_y)
-        y_std = extra['dist'].scale
-        loss = losses['loss'].mean()
-        
-        tensorboard_logs = {
-            "val_loss": loss, # This exact key is needed for metrics
-            "val/kl": losses['loss_kl'].mean(),
-            "val/std": y_std.mean(),
-            "val/mse": losses['loss_mse'].mean(),
-        }
-        return {"val_loss": loss, "log": tensorboard_logs}
-
-    # def training_end(self, outputs):
-    #     logs = self.agg_logs(outputs)
-    #     tensorboard_logs_str = {k: f'{v}' for k, v in logs["log"].items()}
-    #     print(f"step train {self.trainer.global_step}, {tensorboard_logs_str}")
-    #     return logs
-
-    def validation_end(self, outputs):
-        if int(self.hparams["vis_i"]) > 0:
-            self.show_image()
-        logs = self.agg_logs(outputs)
-        tensorboard_logs_str = {k: f'{v}' for k, v in logs["log"].items()}
-
-        # agg and print self.train_logs HACK https://github.com/PyTorchLightning/pytorch-lightning/issues/100
-        train_logs = self.agg_logs(self.train_logs)
-        train_logs_str = {k: f"{v}" for k, v in train_logs.items()}
-        self.train_logs = []
-        print(f"step val {self.trainer.global_step}, {tensorboard_logs_str} {train_logs}")
-        return logs
-
-    def show_image(self):        
-        # https://github.com/PytorchLightning/pytorch-lightning/blob/f8d9f8f/pytorch_lightning/core/lightning.py#L293
-        loader = self.val_dataloader()
-        vis_i = min(int(self.hparams["vis_i"]), len(loader.dataset))
-        # print('vis_i', vis_i)
-        if isinstance(self.hparams["vis_i"], str):
-            image = plot_from_loader(loader, self, i=int(vis_i))
-            plt.show()
-        else:
-            image = plot_from_loader_to_tensor(loader, self, i=vis_i)
-            self.logger.experiment.add_image('val/image', image, self.trainer.global_step)
-
-    def agg_logs(self, outputs):
-        if isinstance(outputs, dict):
-            outputs = [outputs]
-        aggs = {}
-        if len(outputs)>0:
-            for j in outputs[0]:
-                if isinstance(outputs[0][j], dict):
-                    # Take mean of sub dicts
-                    keys = outputs[0][j].keys()
-                    aggs[j] = {k: torch.stack([x[j][k] for x in outputs if k in x[j]]).mean().item() for k in keys}
-                else:
-                    # Take mean of numbers
-                    aggs[j] = torch.stack([x[j] for x in outputs if j in x]).mean().item()
-        return aggs
-
-        # # Log hparams with metric, doesn't work
-        # # self.logger.experiment.add_hparams(self.hparams.__dict__, {"avg_val_loss": avg_loss})
-        # if f"{name}_loss" in outputs[0].keys():
-        #     avg_loss = torch.stack([x[f"{name}_loss"] for x in outputs]).mean()
-        #     assert torch.isfinite(avg_loss)
-        # else:
-        #     avg_loss = 0
-        # return {f"avg_{name}_loss": avg_loss, "log": tensorboard_logs, "progress_bar": {}}
-
-    def test_step(self, *args, **kwargs):
-        return self.validation_step(*args, **kwargs)
-
-    def test_end(self, *args, **kwargs):
-        return self.validation_end(*args, **kwargs)
-
-    def configure_optimizers(self):
-        optim = torch.optim.AdamW(self.parameters(), lr=self.hparams["learning_rate"], weight_decay=0)
-        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optim, patience=self.hparams["patience"], verbose=True, min_lr=1e-7)  # note early stopping has patience 3
-        return [optim], [scheduler]
-
-    def _get_cache_dfs(self):
-        if self._dfs is None:
-            df_train, df_val, df_test = get_smartmeter_df()
-            self._dfs = dict(df_train=df_train, df_val=df_val, df_test=df_test)
-        return self._dfs
-
-    def train_dataloader(self):
-        df_train = self._get_cache_dfs()['df_train']
-        data_train = SmartMeterDataSet(
-            df_train, self.hparams["num_context"], self.hparams["num_extra_target"]
-        )
-        return torch.utils.data.DataLoader(
-            data_train,
-            batch_size=self.hparams["batch_size"],
-            shuffle=True,
-            collate_fn=collate_fns(
-                self.hparams["num_context"], self.hparams["num_extra_target"], sample=True, context_in_target=self.hparams["context_in_target"]
-            ),
-            num_workers=self.hparams["num_workers"],
-        )
-
-    def val_dataloader(self):
-        df_test = self._get_cache_dfs()['df_val']
-        data_test = SmartMeterDataSet(
-            df_test, self.hparams["num_context"], self.hparams["num_extra_target"]
-        )
-        return torch.utils.data.DataLoader(
-            data_test,
-            batch_size=self.hparams["batch_size"],
-            shuffle=False,
-            collate_fn=collate_fns(
-                self.hparams["num_context"], self.hparams["num_extra_target"], sample=False, context_in_target=self.hparams["context_in_target"]
-            ),
-        )
-
-    def test_dataloader(self):
-        df_test = self._get_cache_dfs()['df_test']
-        data_test = SmartMeterDataSet(
-            df_test, self.hparams["num_context"], self.hparams["num_extra_target"]
-        )
-        return torch.utils.data.DataLoader(
-            data_test,
-            batch_size=self.hparams["batch_size"],
-            shuffle=False,
-            collate_fn=collate_fns(
-                self.hparams["num_context"], self.hparams["num_extra_target"], sample=False, context_in_target=self.hparams["context_in_target"]
-            ),
-        )
-
-    @staticmethod
-    def add_suggest(trial):        
-        trial.suggest_loguniform("learning_rate", 1e-5, 1e-2)
-
-        trial.suggest_categorical("hidden_dim", [8*2**i for i in range(8)])
-        trial.suggest_categorical("latent_dim", [8*2**i for i in range(8)])
-        
-        trial.suggest_int("attention_layers", 1, 4)
-        trial.suggest_categorical("n_latent_encoder_layers", [1, 2, 4, 6, 8, 12])
-        trial.suggest_categorical("n_det_encoder_layers", [1, 2, 4, 6, 8, 12])
-        trial.suggest_categorical("n_decoder_layers", [1, 2, 4, 6, 8, 12])
-        trial.suggest_int("num_heads", 8, 8)
-
-        trial.suggest_uniform("dropout", 0, 0.9)
-        trial.suggest_uniform("attention_dropout", 0, 0.9)
-
-        trial.suggest_categorical(
-            "latent_enc_self_attn_type", ['uniform', 'multihead', 'ptmultihead']
-        )
-        trial.suggest_categorical("det_enc_self_attn_type",  ['uniform', 'multihead', 'ptmultihead'])
-        trial.suggest_categorical("det_enc_cross_attn_type", ['uniform', 'multihead', 'ptmultihead'])
-
-        trial.suggest_categorical("batchnorm", [False, True])
-        trial.suggest_categorical("use_self_attn", [False, True])
-        trial.suggest_categorical("use_lvar", [False, True])
-        trial.suggest_categorical("use_deterministic_path", [False, True])
-        trial.suggest_categorical("use_rnn", [True, False])
-
-        trial._user_attrs = {
-            'batch_size': 16,
-            'grad_clip': 40,
-            'max_nb_epochs': 200,
-            'num_workers': 4,
-            'num_context': 24* 4,
-            'vis_i': '670',
-            'num_extra_target': 24*4,
-            'x_dim': 18,
-            'context_in_target': True,
-            'y_dim': 1,
-            'patience': 3,
-            'min_std': 0.005,
-        }        
-        return trial
-
-

src/models/lstm.py

@@ -1,299 +0,0 @@
-import os
-import numpy as np
-import pandas as pd
-import torch
-from tqdm.auto import tqdm
-from torch import nn
-from torch.nn import functional as F
-from torch.utils.data import DataLoader
-from torchvision.datasets import MNIST
-from test_tube import Experiment, HyperOptArgumentParser
-import torchvision.transforms as transforms
-from argparse import ArgumentParser
-import json
-import pytorch_lightning as pl
-from matplotlib import pyplot as plt
-import torch
-import io
-import PIL
-from torchvision.transforms import ToTensor
-
-from src.data.smart_meter import get_smartmeter_df
-
-from src.utils import ObjectDict
-
-
-class SequenceDfDataSet(torch.utils.data.Dataset):
-    def __init__(self, df, hparams, label_names=None, train=True, transforms=None):
-        super().__init__()
-        self.data = df
-        self.hparams = hparams
-        self.label_names = label_names
-        self.train = train
-        self.transforms = transforms
-
-    def __len__(self):
-        return len(self.data) - self.hparams.window_length - self.hparams.target_length - 1
-
-    def iloc(self, idx):
-        k = idx + self.hparams.window_length + self.hparams.target_length
-        j = k - self.hparams.target_length
-        i = j - self.hparams.window_length
-        assert i >= 0
-        assert idx <= len(self.data)
-
-        x_rows = self.data.iloc[i:k].copy()
-        # x_rows = x_rows.drop(columns=self.label_names)
-        # Note the NP models do have access to the previous labels for the context, we will allow the LSTM to do the same. Although it will likely just return an autoregressive solution for the first half...
-        x_rows.loc[x_rows.index[self.hparams.window_length:], self.label_names] = 0
-        assert len(x_rows.loc[x_rows.index[self.hparams.window_length:], self.label_names])>0
-        assert (x_rows.loc[x_rows.index[self.hparams.window_length:], self.label_names]==0).all().all()
-
-        y_rows = self.data[self.label_names].iloc[i+1:k+1].copy()
-        #         print(i,j,k)
-
-        # add seconds since start of window index
-        x_rows["tstp"] = (
-            x_rows["tstp"] - x_rows["tstp"].iloc[0]
-        ).dt.total_seconds() / 86400.0
-        return x_rows, y_rows
-
-    def __getitem__(self, idx):
-        x_rows, y_rows = self.iloc(idx)
-
-        x = x_rows.astype(np.float32).values
-        y = y_rows[self.label_names].astype(np.float32).values
-        return (
-            self.transforms(x).squeeze(0).float(),
-            self.transforms(y).squeeze(0).squeeze(-1).float(),
-        )
-
-
-class LSTMNet(nn.Module):
-    def __init__(self, hparams):
-        super().__init__()
-        self.hparams = hparams
-
-        self.lstm1 = nn.LSTM(
-            input_size=self.hparams.input_size,
-            hidden_size=self.hparams.hidden_size,
-            batch_first=True,
-            num_layers=self.hparams.lstm_layers,
-            bidirectional=self.hparams.bidirectional,
-            dropout=self.hparams.lstm_dropout,
-        )
-        self.hidden_out_size = (
-            self.hparams.hidden_size
-            * (self.hparams.bidirectional + 1)
-        )
-        self.linear = nn.Linear(self.hidden_out_size, 1)
-
-    def forward(self, x):
-        outputs, (h_out, _) = self.lstm1(x)
-        # outputs: [B, T, num_direction * H]
-        y = self.linear(outputs).squeeze(2)
-        return y
-
-
-class LSTM_PL(pl.LightningModule):
-    def __init__(self, hparams):
-        # TODO make label name configurable
-        # TODO make data source configurable
-        super().__init__()
-        self.hparams = ObjectDict()
-        self.hparams.update(
-            hparams.__dict__ if hasattr(hparams, "__dict__") else hparams
-        )
-        self._model = LSTMNet(self.hparams)
-        self._dfs = None
-
-    def forward(self, x):
-        return self._model(x)
-
-    def training_step(self, batch, batch_idx):
-        # REQUIRED
-        x, y = batch
-        y_hat = self.forward(x)
-        y = y[:, self.hparams.window_length:]
-        y_hat = y_hat[:, self.hparams.window_length:]
-        loss = F.mse_loss(y_hat, y)
-        tensorboard_logs = {"train_loss": loss}
-        return {"loss": loss, "log": tensorboard_logs}
-
-    def validation_step(self, batch, batch_idx):
-        x, y = batch
-        y_hat = self.forward(x)
-        y = y[:, self.hparams.window_length:]
-        y_hat = y_hat[:, self.hparams.window_length:]
-        loss = F.mse_loss(y_hat, y)
-        tensorboard_logs = {"val_loss": loss}
-        return {"val_loss": loss, "log": tensorboard_logs}
-
-    def validation_end(self, outputs):
-        # TODO send an image to tensroboard, like in the lighting_anp.py file
-        if int(self.hparams["vis_i"]) > 0:
-            loader = self.val_dataloader()
-            vis_i = min(int(self.hparams["vis_i"]), len(loader.dataset))
-        if isinstance(self.hparams["vis_i"], str):
-            image = plot_from_loader(loader, self, vis_i=vis_i, window_len=self.hparams["window_length"])
-            plt.show()
-        else:
-            image = plot_from_loader_to_tensor(loader, self, vis_i=vis_i, window_len=self.hparams["window_length"])
-            self.logger.experiment.add_image(
-                "val/image", image, self.trainer.global_step
-            )
-
-        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
-        keys = outputs[0]["log"].keys()
-        tensorboard_logs = {
-            k: torch.stack([x["log"][k] for x in outputs if k in x["log"]]).mean()
-            for k in keys
-        }
-        tensorboard_logs_str = {k: f"{v}" for k, v in tensorboard_logs.items()}
-        print(f"step {self.trainer.global_step}, {tensorboard_logs_str}")
-        assert torch.isfinite(avg_loss)
-        return {"avg_val_loss": avg_loss, "log": tensorboard_logs}
-
-    def test_step(self, *args, **kwargs):
-        return self.validation_step(*args, **kwargs)
-
-    def test_end(self, *args, **kwargs):
-        return self.validation_end(*args, **kwargs)
-
-    def configure_optimizers(self):
-        optim = torch.optim.Adam(self.parameters(), lr=self.hparams["learning_rate"])
-        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
-            optim, patience=self.hparams["patience"], verbose=True, min_lr=1e-5
-        )  # note early stopping has patient 3
-        return [optim], [scheduler]
-
-    def _get_cache_dfs(self):
-        if self._dfs is None:
-            df_train, df_val, df_test = get_smartmeter_df()
-            self._dfs = dict(df_train=df_train, df_val=df_val, df_test=df_test)
-        return self._dfs
-
-    @pl.data_loader
-    def train_dataloader(self):
-        df_train = self._get_cache_dfs()["df_train"]
-        dset_train = SequenceDfDataSet(
-            df_train,
-            self.hparams,
-            label_names=["energy(kWh/hh)"],
-            transforms=transforms.ToTensor(),
-            train=True,
-        )
-        return DataLoader(
-            dset_train,
-            batch_size=self.hparams.batch_size,
-            shuffle=True,
-            num_workers=self.hparams.num_workers,
-        )
-
-    @pl.data_loader
-    def val_dataloader(self):
-        df_test = self._get_cache_dfs()["df_val"]
-        dset_test = SequenceDfDataSet(
-            df_test,
-            self.hparams,
-            label_names=["energy(kWh/hh)"],
-            train=False,
-            transforms=transforms.ToTensor(),
-        )
-        return DataLoader(dset_test, batch_size=self.hparams.batch_size, shuffle=False)
-
-    @pl.data_loader
-    def test_dataloader(self):
-        df_test = self._get_cache_dfs()["df_test"]
-        dset_test = SequenceDfDataSet(
-            df_test,
-            self.hparams,
-            label_names=["energy(kWh/hh)"],
-            train=False,
-            transforms=transforms.ToTensor(),
-        )
-        return DataLoader(dset_test, batch_size=self.hparams.batch_size, shuffle=False)
-
-    @staticmethod
-    def add_suggest(trial: optuna.Trial):
-        """
-        Add hyperparam ranges to an optuna trial and typical user attrs.
-        
-        Usage:
-            trial = optuna.trial.FixedTrial(
-                params={         
-                    'hidden_size': 128,
-                }
-            )
-            trial = add_suggest(trial)
-            trainer = pl.Trainer()
-            model = LSTM_PL(dict(**trial.params, **trial.user_attrs), dataset_train,
-                            dataset_test, cache_base_path, norm)
-            trainer.fit(model)
-        """
-        trial.suggest_loguniform("learning_rate", 1e-6, 1e-2)
-        trial.suggest_uniform("lstm_dropout", 0, 0.75)
-        trial.suggest_categorical(
-            "hidden_size", [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]
-        )
-        trial.suggest_categorical("lstm_layers", [1, 2, 3, 4, 6,  8])
-        trial.suggest_categorical("bidirectional", [False, True])
-
-        trial._user_attrs = {
-            "batch_size": 16,
-            "grad_clip": 40,
-            "max_nb_epochs": 200,
-            "num_workers": 4,
-            "vis_i": 670,
-            "input_size": 6,
-            "output_size": 1,
-            "patience": 2,
-        }
-        return trial
-
-
-def plot_from_loader(loader, model, vis_i=670, n=1, window_len=0):
-    dset_test = loader.dataset
-    label_names = dset_test.label_names
-    y_trues = []
-    y_preds = []
-    vis_i = min(vis_i, len(dset_test))
-    for i in tqdm(range(vis_i, vis_i + n)):
-        x_rows, y_rows = dset_test.iloc(i)
-        x, y = dset_test[i]
-        device = next(model.parameters()).device
-        x = x[None, :].to(device)
-        model.eval()
-        with torch.no_grad():
-            y_hat = model.forward(x)
-            y_hat = y_hat.cpu().squeeze(0).numpy()
-
-        dt = y_rows.iloc[0].name
-
-        y_hat_rows = y_rows.copy()
-        y_hat_rows[label_names[0]] = y_hat
-        y_trues.append(y_rows)
-        y_preds.append(y_hat_rows)
-
-    plt.figure()
-    pd.concat(y_trues)[label_names[0]].plot(label="y_true")
-    ylims = plt.ylim()
-    pd.concat(y_preds)[label_names[0]][window_len:].plot(label="y_pred")
-    plt.legend()
-    t_ahead = pd.Timedelta("30T") * model.hparams.target_length
-    plt.title(f"predicting {t_ahead} ahead")
-    plt.ylim(*ylims)
-    # plt.show()
-
-
-def plot_from_loader_to_tensor(*args, **kwargs):
-    plot_from_loader(*args, **kwargs)
-
-    # Send fig to tensorboard
-    buf = io.BytesIO()
-    plt.savefig(buf, format="jpeg")
-    plt.close()
-    buf.seek(0)
-    image = PIL.Image.open(buf)
-    image = ToTensor()(image)  # .unsqueeze(0)
-    return image

src/models/lstm_seqseq.py

@@ -1,281 +0,0 @@
-import os
-import numpy as np
-import pandas as pd
-import torch
-from tqdm.auto import tqdm
-from torch import nn
-from torch.nn import functional as F
-from torch.utils.data import DataLoader
-from torchvision.datasets import MNIST
-from test_tube import Experiment, HyperOptArgumentParser
-from src.data.smart_meter import collate_fns, SmartMeterDataSet, get_smartmeter_df
-import torchvision.transforms as transforms
-from src.plot import plot_from_loader_to_tensor, plot_from_loader
-from argparse import ArgumentParser
-import json
-import pytorch_lightning as pl
-import math
-from matplotlib import pyplot as plt
-import torch
-import io
-import PIL
-from torchvision.transforms import ToTensor
-from src.models.modules import BatchNormSequence
-from src.data.smart_meter import get_smartmeter_df
-
-from src.utils import ObjectDict
-
-def log_prob_sigma(value, loc, log_scale):
-    """A slightly more stable (not confirmed yet) log prob taking in log_var instead of scale.
-    modified from https://github.com/pytorch/pytorch/blob/2431eac7c011afe42d4c22b8b3f46dedae65e7c0/torch/distributions/normal.py#L65
-    """
-    var = torch.exp(log_scale * 2)
-    return (
-        -((value - loc) ** 2) / (2 * var) - log_scale - math.log(math.sqrt(2 * math.pi))
-    )
-
-
-class Seq2SeqNet(nn.Module):
-    def __init__(self, hparams, _min_std = 0.05):
-        super().__init__()
-        self.hparams = hparams
-        self._min_std = _min_std
-
-
-
-        self.norm_input = BatchNormSequence(self.hparams.input_size)
-        self.encoder = nn.LSTM(
-            input_size=self.hparams.input_size,
-            hidden_size=self.hparams.hidden_size,
-            batch_first=True,
-            num_layers=self.hparams.lstm_layers,
-            bidirectional=self.hparams.bidirectional,
-            dropout=self.hparams.lstm_dropout,
-        )
-        self.multihead_attn = nn.MultiheadAttention(self.hparams.hidden_size, num_heads=8)
-
-        self.norm_target = BatchNormSequence(self.hparams.input_size_decoder)
-        self.decoder = nn.LSTM(
-            input_size=self.hparams.input_size_decoder,
-            hidden_size=self.hparams.hidden_size,
-            batch_first=True,
-            num_layers=self.hparams.lstm_layers,
-            bidirectional=self.hparams.bidirectional,
-            dropout=self.hparams.lstm_dropout,
-        )
-        self.hidden_out_size = (
-            self.hparams.hidden_size
-            * (self.hparams.bidirectional + 1)
-        )
-        self.mean = nn.Linear(self.hidden_out_size, self.hparams.output_size)
-        self.std = nn.Linear(self.hidden_out_size, self.hparams.output_size)
-        self._use_lvar = False
-
-    def forward(self, context_x, context_y, target_x, target_y=None):
-        x = torch.cat([context_x, context_y], -1)
-
-        # Sometimes input normalisation can be important, an initial batch norm is a nice way to ensure this
-        x = self.norm_input(x)
-        target_x = self.norm_target(target_x)
-
-        _, (h_out, cell) = self.encoder(x)
-        # hidden = [batch size, n layers * n directions, hid dim]
-        # cell = [batch size, n layers * n directions, hid dim]
-
-        # context_x, d_encoded, target_x = k, v, q
-
-        # query, key, value = target_x, context_x, d_encoded
-        attn_output, _ = self.multihead_attn(h_out.permute(1, 0, 2), h_out.permute(1, 0, 2), h_out.permute(1, 0, 2))
-        h_out = attn_output.permute(1, 0, 2).contiguous()
-        attn_output, _ = self.multihead_attn(cell.permute(1, 0, 2), cell.permute(1, 0, 2), cell.permute(1, 0, 2))
-        cell = attn_output.permute(1, 0, 2).contiguous()
-
-        outputs, (_, _) = self.decoder(target_x, (h_out, cell))
-        # output = [batch size, seq len, hid dim * n directions]
-        
-        # outputs: [B, T, num_direction * H]
-        mean = self.mean(outputs)
-        log_sigma = self.std(outputs)
-        if self._use_lvar:
-            log_sigma = torch.clamp(log_sigma, math.log(self._min_std), -math.log(self._min_std))
-            sigma = torch.exp(log_sigma)
-        else:
-            sigma = self._min_std + (1 - self._min_std) * F.softplus(log_sigma)
-        y_dist=torch.distributions.Normal(mean, sigma)
-        
-        # Loss
-        loss_mse = loss_p = None
-        if target_y is not None:
-            loss_mse = F.mse_loss(mean, target_y, reduction='none')
-            if self._use_lvar:
-                loss_p = -log_prob_sigma(target_y, mean, log_sigma)
-            else:
-                loss_p = -y_dist.log_prob(target_y).mean(-1)
-            
-            if self.hparams["context_in_target"]:
-                loss_p[:context_x.size(1)] /= 100
-                loss_mse[:context_x.size(1)] /= 100
-            # # Don't catch loss on context window
-            # mean = mean[:, self.hparams.num_context:]
-            # log_sigma = log_sigma[:, self.hparams.num_context:]
-
-        y_pred = y_dist.rsample if self.training else y_dist.loc
-        return y_pred, dict(loss_p=loss_p.mean(), loss_mse=loss_mse.mean()), dict(log_sigma=log_sigma, dist=y_dist)
-
-
-class LSTMSeq2Seq_PL(pl.LightningModule):
-    def __init__(self, hparams):
-        # TODO make label name configurable
-        # TODO make data source configurable
-        super().__init__()
-        self.hparams = ObjectDict()
-        self.hparams.update(
-            hparams.__dict__ if hasattr(hparams, "__dict__") else hparams
-        )
-        self.model = Seq2SeqNet(self.hparams)
-        self._dfs = None
-
-    def forward(self, context_x, context_y, target_x, target_y):
-        return self.model(context_x, context_y, target_x, target_y)
-
-    def training_step(self, batch, batch_idx):
-        # REQUIRED
-        assert all(torch.isfinite(d).all() for d in batch)
-        context_x, context_y, target_x, target_y = batch
-        y_dist, losses, extra = self.forward(context_x, context_y, target_x, target_y)
-        loss = losses['loss_p'] # + loss_mse
-        tensorboard_logs = {
-            "train/loss": loss,
-            'train/loss_mse': losses['loss_mse'],
-            "train/loss_p": losses['loss_p'],
-            "train/sigma": torch.exp(extra['log_sigma']).mean()}
-        return {"loss": loss, "log": tensorboard_logs}
-
-    def validation_step(self, batch, batch_idx):
-        context_x, context_y, target_x, target_y = batch
-        assert all(torch.isfinite(d).all() for d in batch)
-        y_dist, losses, extra = self.forward(context_x, context_y, target_x, target_y)
-        loss = losses['loss_p'] # + loss_mse
-        tensorboard_logs = {
-            "val_loss": loss,
-            'val/loss_mse': losses['loss_mse'],
-            "val/loss_p": losses['loss_p'],
-            "val/sigma": torch.exp(extra['log_sigma']).mean()}
-        return {"val_loss": loss, "log": tensorboard_logs}
-
-    def validation_end(self, outputs):
-        if int(self.hparams["vis_i"]) > 0:
-            self.show_image()
-
-        avg_loss = torch.stack([x["val_loss"] for x in outputs]).mean()
-        keys = outputs[0]["log"].keys()
-        tensorboard_logs = {
-            k: torch.stack([x["log"][k] for x in outputs if k in x["log"]]).mean()
-            for k in keys
-        }
-        tensorboard_logs_str = {k: f"{v}" for k, v in tensorboard_logs.items()}
-        print(f"step {self.trainer.global_step}, {tensorboard_logs_str}")
-        assert torch.isfinite(avg_loss)
-        return {"avg_val_loss": avg_loss, "log": tensorboard_logs}
-
-
-    def show_image(self):        
-        # https://github.com/PytorchLightning/pytorch-lightning/blob/f8d9f8f/pytorch_lightning/core/lightning.py#L293
-        loader = self.val_dataloader()
-        vis_i = min(int(self.hparams["vis_i"]), len(loader.dataset))
-        # print('vis_i', vis_i)
-        if isinstance(self.hparams["vis_i"], str):
-            image = plot_from_loader(loader, self, i=int(vis_i))
-            plt.show()
-        else:
-            image = plot_from_loader_to_tensor(loader, self, i=vis_i)
-            self.logger.experiment.add_image('val/image', image, self.trainer.global_step)
-
-    def test_step(self, *args, **kwargs):
-        return self.validation_step(*args, **kwargs)
-
-    def test_end(self, *args, **kwargs):
-        return self.validation_end(*args, **kwargs)
-
-    def configure_optimizers(self):
-        optim = torch.optim.Adam(self.parameters(), lr=self.hparams["learning_rate"])
-        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
-            optim, patience=self.hparams["patience"], verbose=True, min_lr=1e-5
-        )  # note early stopping has patient 3
-        return [optim], [scheduler]
-
-    def _get_cache_dfs(self):
-        if self._dfs is None:
-            df_train, df_val, df_test = get_smartmeter_df()
-            self._dfs = dict(df_train=df_train, df_val=df_val, df_test=df_test)
-        return self._dfs
-
-    @pl.data_loader
-    def train_dataloader(self):
-        df_train = self._get_cache_dfs()['df_train']
-        data_train = SmartMeterDataSet(
-            df_train, self.hparams["num_context"], self.hparams["num_extra_target"]
-        )
-        return torch.utils.data.DataLoader(
-            data_train,
-            batch_size=self.hparams["batch_size"],
-            shuffle=True,
-            collate_fn=collate_fns(
-                self.hparams["num_context"], self.hparams["num_extra_target"], sample=True, context_in_target=self.hparams["context_in_target"]
-            ),
-            num_workers=self.hparams["num_workers"],
-        )
-
-    @pl.data_loader
-    def val_dataloader(self):
-        df_test = self._get_cache_dfs()['df_val']
-        data_test = SmartMeterDataSet(
-            df_test, self.hparams["num_context"], self.hparams["num_extra_target"]
-        )
-        return torch.utils.data.DataLoader(
-            data_test,
-            batch_size=self.hparams["batch_size"],
-            shuffle=False,
-            collate_fn=collate_fns(
-                self.hparams["num_context"], self.hparams["num_extra_target"], sample=False, context_in_target=self.hparams["context_in_target"]
-            ),
-        )
-
-    @pl.data_loader
-    def test_dataloader(self):
-        df_test = self._get_cache_dfs()['df_test']
-        data_test = SmartMeterDataSet(
-            df_test, self.hparams["num_context"], self.hparams["num_extra_target"]
-        )
-        return torch.utils.data.DataLoader(
-            data_test,
-            batch_size=self.hparams["batch_size"],
-            shuffle=False,
-            collate_fn=collate_fns(
-                self.hparams["num_context"], self.hparams["num_extra_target"], sample=False, context_in_target=self.hparams["context_in_target"]
-            ),
-        )
-
-    @staticmethod
-    def add_suggest(trial):
-        trial.suggest_loguniform("learning_rate", 1e-5, 1e-2)
-        trial.suggest_uniform("lstm_dropout", 0, 0.75)
-        trial.suggest_categorical("hidden_size", [1, 2, 4, 8, 16, 32, 64, 128, 256, 512])    
-        trial.suggest_categorical("lstm_layers", [1, 2, 4, 8])    
-        trial.suggest_categorical("bidirectional", [False, True])    
-        
-
-        trial._user_attrs = {
-                'batch_size': 16,
-                'grad_clip': 40,
-                'max_nb_epochs': 200,
-                'num_workers': 4,
-                'num_extra_target': 24*4,
-                'vis_i': '670',
-                'num_context': 24*4,
-                'input_size': 18,
-                'input_size_decoder': 17,
-                'context_in_target': True,
-                'output_size': 1
-        }
-        return trial

src/models/modules.py

@@ -1,410 +0,0 @@
-import torch
-from torch import nn
-import torch.nn.functional as F
-import math
-import numpy as np
-# from .attention import Attention as PtAttention
-
-
-class LSTMBlock(nn.Module):
-    def __init__(
-        self, in_channels, out_channels, dropout=0, batchnorm=False, bias=False, num_layers=1
-    ):
-        super().__init__()
-        self._lstm = nn.LSTM(
-                input_size=in_channels,
-                hidden_size=out_channels,
-                num_layers=num_layers,
-                dropout=dropout,
-                batch_first=True,
-                bias=bias
-        )
-
-    def forward(self, x):
-        return self._lstm(x)[0]
-
-
-class BatchNormSequence(nn.Module):
-    """Applies batch norm on features of a batch first sequence."""
-    def __init__(
-        self, out_channels
-    ):
-        super().__init__()
-        self.norm = nn.BatchNorm1d(out_channels)
-
-    def forward(self, x):
-        # x.shape is (Batch, Sequence, Channels)
-        # Now we want to apply batchnorm and dropout to the channels. So we put it in shape
-        # (Batch, Channels, Sequence) so we can use BatchNorm1d
-        x = x.permute(0, 2, 1)
-        x = self.norm(x)
-        return x.permute(0, 2, 1)
-
-class NPBlockRelu2d(nn.Module):
-    """Block for Neural Processes."""
-
-    def __init__(
-        self, in_channels, out_channels, dropout=0, batchnorm=False, bias=False
-    ):
-        super().__init__()
-        self.linear = nn.Linear(in_channels, out_channels, bias=bias)
-        self.act = nn.ReLU()
-        self.dropout = nn.Dropout2d(dropout)
-        self.norm = nn.BatchNorm2d(out_channels) if batchnorm else False
-
-    def forward(self, x):
-        # x.shape is (Batch, Sequence, Channels)
-        # We pass a linear over it which operates on the Channels
-        x = self.act(self.linear(x))
-
-        # Now we want to apply batchnorm and dropout to the channels. So we put it in shape
-        # (Batch, Channels, Sequence, None) so we can use Dropout2d & BatchNorm2d
-        x = x.permute(0, 2, 1)[:, :, :, None]
-
-        if self.norm:
-            x = self.norm(x)
-
-        x = self.dropout(x)
-        return x[:, :, :, 0].permute(0, 2, 1)
-
-
-class BatchMLP(nn.Module):
-    """Apply MLP to the final axis of a 3D tensor (reusing already defined MLPs).
-
-    Args:
-        input: input tensor of shape [B,n,d_in].
-        output_sizes: An iterable containing the output sizes of the MLP as defined 
-            in `basic.Linear`.
-    Returns:
-        tensor of shape [B,n,d_out] where d_out=output_size
-    """
-
-    def __init__(
-        self, input_size, output_size, num_layers=2, dropout=0, batchnorm=False
-    ):
-        super().__init__()
-        self.input_size = input_size
-        self.output_size = output_size
-        self.num_layers = num_layers
-
-        self.initial = NPBlockRelu2d(
-            input_size, output_size, dropout=dropout, batchnorm=batchnorm
-        )
-        self.encoder = nn.Sequential(
-            *[
-                NPBlockRelu2d(
-                    output_size, output_size, dropout=dropout, batchnorm=batchnorm
-                )
-                for _ in range(num_layers - 2)
-            ]
-        )
-        self.final = nn.Linear(output_size, output_size)
-
-    def forward(self, x):
-        x = self.initial(x)
-        x = self.encoder(x)
-        return self.final(x)
-
-
-class AttnLinear(nn.Module):
-    def __init__(self, in_channels, out_channels):
-        super().__init__()
-        self.linear = nn.Linear(in_channels, out_channels, bias=False)
-        torch.nn.init.normal_(self.linear.weight, std=in_channels ** -0.5)
-
-    def forward(self, x):
-        x = self.linear(x)
-        return x
-
-
-class Attention(nn.Module):
-    def __init__(
-        self,
-        hidden_dim,
-        attention_type,
-        attention_layers=2,
-        n_heads=8,
-        x_dim=1,
-        rep="mlp",
-        dropout=0,
-        batchnorm=False,
-    ):
-        super().__init__()
-        self._rep = rep
-
-        if self._rep == "mlp":
-            self.batch_mlp_k = BatchMLP(
-                x_dim,
-                hidden_dim,
-                attention_layers,
-                dropout=dropout,
-                batchnorm=batchnorm,
-            )
-            self.batch_mlp_q = BatchMLP(
-                x_dim,
-                hidden_dim,
-                attention_layers,
-                dropout=dropout,
-                batchnorm=batchnorm,
-            )
-
-        if attention_type == "uniform":
-            self._attention_func = self._uniform_attention
-        elif attention_type == "laplace":
-            self._attention_func = self._laplace_attention
-        elif attention_type == "dot":
-            self._attention_func = self._dot_attention
-        elif attention_type == "multihead":
-            self._W_k = nn.ModuleList(
-                [AttnLinear(hidden_dim, hidden_dim) for _ in range(n_heads)]
-            )
-            self._W_v = nn.ModuleList(
-                [AttnLinear(hidden_dim, hidden_dim) for _ in range(n_heads)]
-            )
-            self._W_q = nn.ModuleList(
-                [AttnLinear(hidden_dim, hidden_dim) for _ in range(n_heads)]
-            )
-            self._W = AttnLinear(n_heads * hidden_dim, hidden_dim)
-            self._attention_func = self._multihead_attention
-            self.n_heads = n_heads
-        elif attention_type == "ptmultihead":
-            self._W = torch.nn.MultiheadAttention(
-                hidden_dim, n_heads, bias=False, dropout=dropout
-            )
-            self._attention_func = self._pytorch_multihead_attention
-        else:
-            raise NotImplementedError
-
-    def forward(self, k, v, q):
-        if self._rep == "mlp":
-            k = self.batch_mlp_k(k)
-            q = self.batch_mlp_q(q)
-        rep = self._attention_func(k, v, q)
-        return rep
-
-    def _uniform_attention(self, k, v, q):
-        total_points = q.shape[1]
-        rep = torch.mean(v, dim=1, keepdim=True)
-        rep = rep.repeat(1, total_points, 1)
-        return rep
-
-    def _laplace_attention(self, k, v, q, scale=0.5):
-        k_ = k.unsqueeze(1)
-        v_ = v.unsqueeze(2)
-        unnorm_weights = torch.abs((k_ - v_) * scale)
-        unnorm_weights = unnorm_weights.sum(dim=-1)
-        weights = torch.softmax(unnorm_weights, dim=-1)
-        rep = torch.einsum("bik,bkj->bij", weights, v)
-        return rep
-
-    def _dot_attention(self, k, v, q):
-        scale = q.shape[-1] ** 0.5
-        unnorm_weights = torch.einsum("bjk,bik->bij", k, q) / scale
-        weights = torch.softmax(unnorm_weights, dim=-1)
-
-        rep = torch.einsum("bik,bkj->bij", weights, v)
-        return rep
-
-    def _multihead_attention(self, k, v, q):
-        outs = []
-        for i in range(self.n_heads):
-            k_ = self._W_k[i](k)
-            v_ = self._W_v[i](v)
-            q_ = self._W_q[i](q)
-            out = self._dot_attention(k_, v_, q_)
-            outs.append(out)
-        outs = torch.stack(outs, dim=-1)
-        outs = outs.view(outs.shape[0], outs.shape[1], -1)
-        rep = self._W(outs)
-        return rep
-
-    def _pytorch_multihead_attention(self, k, v, q):
-        # Pytorch multiheaded attention takes inputs if diff order and permutation
-        q = q.permute(1, 0, 2)
-        k = k.permute(1, 0, 2)
-        v = v.permute(1, 0, 2)
-        o = self._W(q, k, v)[0]
-        return o.permute(1, 0, 2)
-
-
-class LatentEncoder(nn.Module):
-    def __init__(
-        self,
-        input_dim,
-        hidden_dim=32,
-        latent_dim=32,
-        self_attention_type="dot",
-        n_encoder_layers=3,
-        min_std=0.01,
-        batchnorm=False,
-        dropout=0,
-        attention_dropout=0,
-        use_lvar=False,
-        use_self_attn=False,
-        attention_layers=2,
-        use_lstm=False
-    ):
-        super().__init__()
-        # self._input_layer = nn.Linear(input_dim, hidden_dim)
-        if use_lstm:
-            self._encoder = LSTMBlock(input_dim, hidden_dim, batchnorm=batchnorm, dropout=dropout, num_layers=n_encoder_layers)
-        else:
-            self._encoder = BatchMLP(input_dim, hidden_dim, batchnorm=batchnorm, dropout=dropout, num_layers=n_encoder_layers)
-        if use_self_attn:
-            self._self_attention = Attention(
-                hidden_dim,
-                self_attention_type,
-                attention_layers,
-                rep="identity",
-                dropout=attention_dropout,
-            )
-        self._penultimate_layer = nn.Linear(hidden_dim, hidden_dim)
-        self._mean = nn.Linear(hidden_dim, latent_dim)
-        self._log_var = nn.Linear(hidden_dim, latent_dim)
-        self._min_std = min_std
-        self._use_lvar = use_lvar
-        self._use_lstm = use_lstm
-        self._use_self_attn = use_self_attn
-
-    def forward(self, x, y):
-        encoder_input = torch.cat([x, y], dim=-1)
-
-        # Pass final axis through MLP
-        encoded = self._encoder(encoder_input)
-
-        # Aggregator: take the mean over all points
-        if self._use_self_attn:
-            attention_output = self._self_attention(encoded, encoded, encoded)
-            mean_repr = attention_output.mean(dim=1)
-        else:
-            mean_repr = encoded.mean(dim=1)
-
-        # Have further MLP layers that map to the parameters of the Gaussian latent
-        mean_repr = torch.relu(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)
-
-        if self._use_lvar:
-            # Clip it in the log domain, so it can only approach self.min_std, this helps avoid mode collapase
-            # 2 ways, a better but untested way using the more stable log domain, and the way from the deepmind repo
-            log_var = F.logsigmoid(log_var)
-            log_var = torch.clamp(log_var, np.log(self._min_std), -np.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):
-    def __init__(
-        self,
-        input_dim,
-        x_dim,
-        hidden_dim=32,
-        n_d_encoder_layers=3,
-        self_attention_type="dot",
-        cross_attention_type="dot",
-        use_self_attn=False,
-        attention_layers=2,
-        batchnorm=False,
-        dropout=0,
-        attention_dropout=0,
-        use_lstm=False,
-    ):
-        super().__init__()
-        self._use_self_attn = use_self_attn
-        # self._input_layer = nn.Linear(input_dim, hidden_dim)
-        if use_lstm:
-            self._d_encoder = LSTMBlock(input_dim, hidden_dim, batchnorm=batchnorm, dropout=dropout, num_layers=n_d_encoder_layers)
-        else:
-            self._d_encoder = BatchMLP(input_dim, hidden_dim, batchnorm=batchnorm, dropout=dropout, num_layers=n_d_encoder_layers)
-        if use_self_attn:
-            self._self_attention = Attention(
-                hidden_dim,
-                self_attention_type,
-                attention_layers,
-                rep="identity",
-                dropout=attention_dropout,
-            )
-        self._cross_attention = Attention(
-            hidden_dim,
-            cross_attention_type,
-            x_dim=x_dim,
-            attention_layers=attention_layers,
-        )
-
-    def forward(self, context_x, context_y, target_x):
-        # Concatenate x and y along the filter axes
-        d_encoder_input = torch.cat([context_x, context_y], dim=-1)
-
-        # Pass final axis through MLP
-        d_encoded = self._d_encoder(d_encoder_input)
-
-        if self._use_self_attn:
-            d_encoded = self._self_attention(d_encoded, d_encoded, d_encoded)
-
-        # Apply attention as mean aggregation
-        h = self._cross_attention(context_x, d_encoded, target_x)
-
-        return h
-
-
-class Decoder(nn.Module):
-    def __init__(
-        self,
-        x_dim,
-        y_dim,
-        hidden_dim=32,
-        latent_dim=32,
-        n_decoder_layers=3,
-        use_deterministic_path=True,
-        min_std=0.01,
-        use_lvar=False,
-        batchnorm=False,
-        dropout=0,
-        use_lstm=False,
-    ):
-        super(Decoder, self).__init__()
-        self._target_transform = nn.Linear(x_dim, hidden_dim)
-        if use_deterministic_path:
-            hidden_dim_2 = 2 * hidden_dim + latent_dim
-        else:
-            hidden_dim_2 = hidden_dim + latent_dim
-            
-        if use_lstm:
-            self._decoder = LSTMBlock(hidden_dim_2, hidden_dim_2, batchnorm=batchnorm, dropout=dropout, num_layers=n_decoder_layers)
-        else:
-            self._decoder = BatchMLP(hidden_dim_2, hidden_dim_2, batchnorm=batchnorm, dropout=dropout, num_layers=n_decoder_layers)
-        self._mean = nn.Linear(hidden_dim_2, y_dim)
-        self._std = nn.Linear(hidden_dim_2, y_dim)
-        self._use_deterministic_path = use_deterministic_path
-        self._min_std = min_std
-        self._use_lvar = use_lvar
-
-    def forward(self, r, z, target_x):
-        # concatenate target_x and representation
-        x = self._target_transform(target_x)
-
-        if self._use_deterministic_path:
-            z = torch.cat([r, z], dim=-1)
-
-        r = torch.cat([z, x], dim=-1)
-
-        r = self._decoder(r)
-
-        # Get the mean and the variance
-        mean = self._mean(r)
-        log_sigma = self._std(r)
-
-        # Bound or clamp the variance
-        if self._use_lvar:
-            log_sigma = torch.clamp(log_sigma, math.log(self._min_std), -math.log(self._min_std))
-            sigma = torch.exp(log_sigma)
-        else:
-            sigma = self._min_std + (1 - self._min_std) * F.softplus(log_sigma)
-
-        dist = torch.distributions.Normal(mean, sigma)
-        return dist, log_sigma