batchnorm etc

neural_processes/lightning.py

@@ -85,6 +85,7 @@ class PL_Seq2Seq(pl.LightningModule):
         # For test use a -logp only
         loss = -y_dist.log_prob(target_y).mean()
         tensorboard_logs = {"test_" + k: v for k, v in losses.items()}
+        tensorboard_logs["test_score"] = loss
         assert torch.isfinite(loss)
         return {"test_loss": loss, "log": tensorboard_logs}
 

neural_processes/models/lstm_seqseq.py

@@ -42,7 +42,10 @@ class Seq2SeqNet(nn.Module):
         self.hparams = hparams
         self._min_std = _min_std
 
-        self.norm_input = BatchNormSequence(self.hparams.input_size)
+        # Sometimes input normalisation can be important, an initial batch norm is a nice way to ensure this https://stackoverflow.com/a/46772183/221742
+        self.norm_input = BatchNormSequence(self.hparams.input_size, affine=False)
+        self.norm_target = BatchNormSequence(self.hparams.input_size_decoder, affine=False)
+
         self.encoder = nn.LSTM(
             input_size=self.hparams.input_size,
             hidden_size=self.hparams.hidden_size,
@@ -58,7 +61,6 @@ class Seq2SeqNet(nn.Module):
             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,

neural_processes/models/neural_process/lightning.py

@@ -16,24 +16,25 @@ class PL_NeuralProcess(PL_Seq2Seq):
         MODEL_CLS=MODEL_CLS, **kwargs)
 
     DEFAULT_ARGS = {
-        'dropout': 0,
-        'learning_rate': 0.006,
-        'attention_dropout': 0,
+        'dropout': 0.1,
+        'learning_rate': 0.004,
+        'attention_dropout': 0.5,
         'batchnorm': False,
         'attention_layers': 2,
         'det_enc_cross_attn_type': 'uniform',
         'det_enc_self_attn_type': 'uniform',
         'latent_enc_self_attn_type': 'uniform',
         'num_heads_power': 3,
-        'hidden_dim_power': 6,
+        'hidden_dim_power': 3,
         'latent_dim_power': 5,
+        'n_latent_encoder_layers': 3,
+        'n_det_encoder_layers': 3,
         'n_decoder_layers': 4,
-        'n_latent_encoder_layers': 2,
-        'use_deterministic_path': False,
-        'n_det_encoder_layers': 4,
+        'use_deterministic_path': True,
         'use_lvar': False,
         'use_self_attn': False,
         'use_rnn': False,
+        'bnorm_inputs': True
     }
 
 

neural_processes/models/neural_process/model.py

@@ -231,9 +231,9 @@ class NeuralProcess(nn.Module):
         self._use_rnn = use_rnn
         self.context_in_target = context_in_target
 
-        # Sometimes input normalisation can be important, an initial batch norm is a nice way to ensure this
-        self.norm_x = BatchNormSequence(x_dim)
-        self.norm_y = BatchNormSequence(y_dim)
+        # Sometimes input normalisation can be important, an initial batch norm is a nice way to ensure this https://stackoverflow.com/a/46772183/221742
+        self.norm_x = BatchNormSequence(x_dim, affine=False)
+        self.norm_y = BatchNormSequence(y_dim, affine=False)
 
         if self._use_rnn:
             self._lstm_x = nn.LSTM(
@@ -309,10 +309,11 @@ class NeuralProcess(nn.Module):
         
         device = next(self.parameters()).device
         
-        # https://stackoverflow.com/a/46772183/221742
-        target_x = self.norm_x(target_x)
-        context_x = self.norm_x(context_x)
-        context_y = self.norm_y(context_y)
+        if self.hparams.get('bnorm_inputs', True):
+            # https://stackoverflow.com/a/46772183/221742
+            target_x = self.norm_x(target_x)
+            context_x = self.norm_x(context_x)
+            context_y = self.norm_y(context_y)
 
         if self._use_rnn:
             # see https://arxiv.org/abs/1910.09323 where x is substituted with h = RNN(x)

neural_processes/models/transformer.py

@@ -22,7 +22,10 @@ class NetTransformer(nn.Module):
 
         hidden_out_size = self.hparams.hidden_out_size
         enc_x_dim = self.hparams.x_dim + self.hparams.y_dim
-        # self.enc_norm = BatchNormSequence(enc_x_dim)
+
+        # Sometimes input normalisation can be important, an initial batch norm is a nice way to ensure this https://stackoverflow.com/a/46772183/221742
+        self.enc_norm = BatchNormSequence(enc_x_dim, affine=False)
+        
         self.enc_emb = nn.Linear(enc_x_dim, hidden_out_size)
         encoder_norm = nn.LayerNorm(hidden_out_size)
         layer_enc = nn.TransformerEncoderLayer(
@@ -54,7 +57,7 @@ class NetTransformer(nn.Module):
         x = x.detach()
         x_key_padding_mask = ~x_mask.any(-1)
 
-        x = self.enc_emb(x).permute(1, 0, 2)
+        x = self.enc_emb(self.enc_norm(x)).permute(1, 0, 2)
         
         outputs = self.encoder(x, src_key_padding_mask=x_key_padding_mask).permute(
             1, 0, 2

neural_processes/models/transformer_seq2seq.py

@@ -43,7 +43,11 @@ class TransformerSeq2SeqNet(nn.Module):
         self._min_std = hparams.min_std
 
         hidden_out_size = self.hparams.hidden_out_size
-        self.enc_norm = BatchNormSequence(self.hparams.input_size)
+
+        # Sometimes input normalisation can be important, an initial batch norm is a nice way to ensure this https://stackoverflow.com/a/46772183/221742
+        self.enc_norm = BatchNormSequence(self.hparams.input_size, affine=False)
+        self.dec_norm = BatchNormSequence(self.hparams.input_size_decoder, affine=False)
+
         self.enc_emb = nn.Linear(self.hparams.input_size, hidden_out_size)
         encoder_norm = nn.LayerNorm(hidden_out_size)
         layer_enc = nn.TransformerEncoderLayer(
@@ -57,7 +61,6 @@ class TransformerSeq2SeqNet(nn.Module):
             layer_enc, num_layers=self.hparams.nlayers, norm=encoder_norm
         )
 
-        self.dec_norm = BatchNormSequence(self.hparams.input_size_decoder)
         self.dec_emb = nn.Linear(self.hparams.input_size_decoder, hidden_out_size)
         layer_dec = nn.TransformerDecoderLayer(
             d_model=hidden_out_size,

neural_processes/modules/attention.py

@@ -136,4 +136,4 @@ class Attention(nn.Module):
 
     def _pytorch_multihead_attention(self, k, v, q):
         # Pytorch multiheaded attention takes inputs if diff order and permutation
-        return batch_first_attention(self._w, q=q, k=k, v=v)[0]
+        return batch_first_attention(self._W, q=q, k=k, v=v)[0]

neural_processes/modules/modules.py

@@ -27,15 +27,15 @@ class LSTMBlock(nn.Module):
 class BatchNormSequence(nn.Module):
     """Applies batch norm on features of a batch first sequence."""
     def __init__(
-        self, out_channels
+        self, out_channels, **kwargs
     ):
         super().__init__()
-        self.norm = nn.BatchNorm1d(out_channels)
+        self.norm = nn.BatchNorm1d(out_channels, **kwargs)
 
     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
+        # (Batch, Channels, Sequence) which is what BatchNorm1d expects
         x = x.permute(0, 2, 1)
         x = self.norm(x)
         return x.permute(0, 2, 1)