diff --git a/neural_processes/lightning.py b/neural_processes/lightning.py index 8b0b948..ee4a940 100644 --- a/neural_processes/lightning.py +++ b/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} diff --git a/neural_processes/models/lstm_seqseq.py b/neural_processes/models/lstm_seqseq.py index 33c0243..48b7005 100644 --- a/neural_processes/models/lstm_seqseq.py +++ b/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, diff --git a/neural_processes/models/neural_process/lightning.py b/neural_processes/models/neural_process/lightning.py index f5075b5..327a1ed 100644 --- a/neural_processes/models/neural_process/lightning.py +++ b/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 } diff --git a/neural_processes/models/neural_process/model.py b/neural_processes/models/neural_process/model.py index f379042..d5605af 100644 --- a/neural_processes/models/neural_process/model.py +++ b/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) diff --git a/neural_processes/models/transformer.py b/neural_processes/models/transformer.py index 0fad302..8846e78 100644 --- a/neural_processes/models/transformer.py +++ b/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 diff --git a/neural_processes/models/transformer_seq2seq.py b/neural_processes/models/transformer_seq2seq.py index e611d9d..c91047e 100644 --- a/neural_processes/models/transformer_seq2seq.py +++ b/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, diff --git a/neural_processes/modules/attention.py b/neural_processes/modules/attention.py index beddccf..7f60446 100644 --- a/neural_processes/modules/attention.py +++ b/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] diff --git a/neural_processes/modules/modules.py b/neural_processes/modules/modules.py index 2d8f788..56332a0 100644 --- a/neural_processes/modules/modules.py +++ b/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)