nicer plots, more classes

seq2seq_time/models/lstm.py

@@ -20,9 +20,11 @@ class LSTM(nn.Module):
 
     def forward(self, past_x, past_y, future_x, future_y=None):
         device = next(self.parameters()).device
-        future_y_fake = (
-            torch.ones(past_y.shape[0], future_x.shape[1], past_y.shape[2]).float().to(device) * self.nan_value
-        )
+        B, S, _ = future_x.shape
+        future_y_fake = past_y[:, -1:, :].repeat(1, S, 1).to(device)
+        # future_y_fake = (
+        #     torch.ones(past_y.shape[0], future_x.shape[1], past_y.shape[2]).float().to(device) * self.nan_value
+        # )
         context = torch.cat([past_x, past_y], -1).detach()
         target = torch.cat([future_x, future_y_fake], -1).detach()
         x = torch.cat([context, target * 1], 1).detach()

seq2seq_time/models/neural_process.py

@@ -32,7 +32,11 @@ class LSTMBlock(nn.Module):
 
 
 class NPBlockRelu2d(nn.Module):
-    """Block for Neural Processes."""
+    """
+    Block for Neural Processes.
+    
+    We want to apply batchnorm and dropout to the channels. We reshape so we can use Dropout2d & BatchNorm2d
+    """
 
     def __init__(
         self, in_channels, out_channels, dropout=0, batchnorm=False, bias=False
@@ -101,7 +105,6 @@ class Attention(nn.Module):
     def __init__(
         self,
         hidden_dim,
-        attention_type,
         attention_layers=2,
         n_heads=8,
         x_dim=1,
@@ -155,48 +158,33 @@ class LatentEncoder(nn.Module):
         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_self_attn=True,
         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._encoder = BatchMLP(
+            input_dim,
+            hidden_dim,
+            batchnorm=batchnorm,
+            dropout=dropout,
+            num_layers=n_encoder_layers,
+        )
+        self._self_attention = Attention(
+            hidden_dim,
+            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_lstm = use_lstm
-        self._use_self_attn = use_self_attn
 
     def forward(self, x, y):
         encoder_input = torch.cat([x, y], dim=-1)
@@ -205,11 +193,8 @@ class LatentEncoder(nn.Module):
         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)
+        attention_output = self._self_attention(encoded, encoded, encoded)
+        mean_repr = attention_output.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))
@@ -230,45 +215,28 @@ class DeterministicEncoder(nn.Module):
         x_dim,
         hidden_dim=32,
         n_d_encoder_layers=3,
-        self_attention_type="dot",
-        cross_attention_type="dot",
-        use_self_attn=True,
         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._d_encoder = BatchMLP(
+            input_dim,
+            hidden_dim,
+            batchnorm=batchnorm,
+            dropout=dropout,
+            num_layers=n_d_encoder_layers,
+        )
+        self._self_attention = Attention(
+            hidden_dim,
+            attention_layers,
+            rep="identity",
+            dropout=attention_dropout,
+        )
         self._cross_attention = Attention(
             hidden_dim,
-            cross_attention_type,
             x_dim=x_dim,
             attention_layers=attention_layers,
         )
@@ -280,8 +248,7 @@ class DeterministicEncoder(nn.Module):
         # 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)
+        d_encoded = self._self_attention(d_encoded, d_encoded, d_encoded)
 
         # Apply attention as mean aggregation
         h = self._cross_attention(past_x, d_encoded, future_x)
@@ -301,7 +268,6 @@ class Decoder(nn.Module):
         min_std=0.01,
         batchnorm=False,
         dropout=0,
-        use_lstm=False,
     ):
         super(Decoder, self).__init__()
         self._future_transform = nn.Linear(x_dim, hidden_dim)
@@ -310,22 +276,14 @@ class Decoder(nn.Module):
         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._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
@@ -363,18 +321,14 @@ class RANP(nn.Module):
         latent_dim=32,  # size of latent space
         n_latent_encoder_layers=2,
         n_det_encoder_layers=2,  # number of deterministic encoder layers
-        n_decoder_layers=2,
+        n_decoder_layers=4,
         use_deterministic_path=True,
         min_std=0.01,  # To avoid collapse use a minimum standard deviation, should be much smaller than variation in labels
         dropout=0,
-        use_self_attn=True,
         attention_dropout=0,
         batchnorm=False,
         attention_layers=2,
         use_rnn=True,  # use RNN/LSTM
-        use_lstm_le=False,  # use another LSTM in latent encoder instead of MLP
-        use_lstm_de=False,  # use another LSTM in determinstic encoder instead of MLP
-        use_lstm_d=False,  # use another lstm in decoder instead of MLP
         **kwargs,
     ):
 
@@ -399,11 +353,9 @@ class RANP(nn.Module):
             n_encoder_layers=n_latent_encoder_layers,
             attention_layers=attention_layers,
             dropout=dropout,
-            use_self_attn=use_self_attn,
             attention_dropout=attention_dropout,
             batchnorm=batchnorm,
             min_std=min_std,
-            use_lstm=use_lstm_le,
         )
 
         self._deterministic_encoder = DeterministicEncoder(
@@ -412,11 +364,9 @@ class RANP(nn.Module):
             hidden_dim=hidden_dim,
             n_d_encoder_layers=n_det_encoder_layers,
             attention_layers=attention_layers,
-            use_self_attn=use_self_attn,
             dropout=dropout,
             batchnorm=batchnorm,
             attention_dropout=attention_dropout,
-            use_lstm=use_lstm_de,
         )
 
         self._decoder = Decoder(
@@ -429,7 +379,6 @@ class RANP(nn.Module):
             min_std=min_std,
             n_decoder_layers=n_decoder_layers,
             use_deterministic_path=use_deterministic_path,
-            use_lstm=use_lstm_d,
         )
         self._use_deterministic_path = use_deterministic_path
 
@@ -443,19 +392,17 @@ class RANP(nn.Module):
             x, _ = self._lstm(x)
             past_x = x[:, :S]
             future_x = x[:, S:]
-            # future_x, _ = self._lstm(future_x)
-            # past_x, _ = self._lstm(past_x)
 
         dist_prior, log_var_prior = self._latent_encoder(past_x, past_y)
 
         if (future_y is not None):
-            dist_post, log_var_post = self._latent_encoder(future_x, future_y)
+            y = torch.cat([past_y, future_y], 1)
+            dist_post, log_var_post = self._latent_encoder(x, y)
 
         if self.training:
             z = dist_prior.rsample()
         else:
             z = dist_prior.loc
-
         num_targets = future_x.size(1)
         z = z.unsqueeze(1).repeat(1, num_targets, 1)  # [B, T_target, H]
 
@@ -478,5 +425,20 @@ class RANP(nn.Module):
                 :, : past_x.size(1)
             ].mean()
             loss = (kl_loss - log_p).mean()
-        return dist, {'loss':loss}
+        return dist, {'loss': loss}
 
+
+# class NP(RANP):
+#     """Recurrent Attentive Neural Process for Sequential Data."""
+#     def __init__(
+#         self,
+#         use_self_attn=True,
+#         # TODO use cross attention flag
+#         use_rnn=True,  # use RNN/LSTM
+#         use_lstm_le=False,  # use another LSTM in latent encoder instead of MLP
+#         use_lstm_de=False,  # use another LSTM in determinstic encoder instead of MLP
+#         use_lstm_d=False,  # use another lstm in decoder instead of MLP
+#         **kwargs,
+#     ):
+#     kwargs
+#     super().__init__(**kwargs)

seq2seq_time/models/transformer.py

@@ -2,12 +2,13 @@ import torch
 from torch import nn
 from torch.nn import functional as F
 
+from ..util import mask_upper_triangular
 
 class Transformer(nn.Module):
     """
     A single transformer, masking nan or 0
     """
-    def __init__(self, x_dim, y_dim, attention_dropout=0, nhead=8, nlayers=2, hidden_size=16, nan_value=0, min_std=0.01):
+    def __init__(self, x_dim, y_dim, attention_dropout=0, nhead=8, nlayers=8, hidden_size=32, nan_value=0, min_std=0.01):
         super().__init__()
         self._min_std = min_std
         self.nan_value = nan_value
@@ -17,7 +18,7 @@ class Transformer(nn.Module):
         encoder_norm = nn.LayerNorm(hidden_size)
         layer_enc = nn.TransformerEncoderLayer(
             d_model=hidden_size,
-            dim_feedforward=hidden_size*4,
+            dim_feedforward=hidden_size*8,
             dropout=attention_dropout,
             nhead=nhead,
             # activation
@@ -30,9 +31,11 @@ class Transformer(nn.Module):
 
     def forward(self, past_x, past_y, future_x, future_y=None):
         device = next(self.parameters()).device
-        future_y_fake = (
-            torch.ones(past_y.shape[0], future_x.shape[1], past_y.shape[2]).float().to(device) * self.nan_value
-        )
+        B, S, _ = future_x.shape
+        future_y_fake = past_y[:, -1:, :].repeat(1, S, 1).to(device)
+        # future_y_fake = (
+        #     torch.ones(past_y.shape[0], future_x.shape[1], past_y.shape[2]).float().to(device) * past_y[:, -1].repeat(B, S, 1)
+        # )
         context = torch.cat([past_x, past_y], -1).detach()
         target = torch.cat([future_x, future_y_fake], -1).detach()
         x = torch.cat([context, target * 1], 1).detach()
@@ -44,8 +47,12 @@ class Transformer(nn.Module):
         x_key_padding_mask = ~x_mask.any(-1)
 
         x = self.enc_emb(x).permute(1, 0, 2)
+
+        B, S, _ = x.shape
+        mask = mask_upper_triangular(S, device)
         
-        outputs = self.encoder(x, src_key_padding_mask=x_key_padding_mask).permute(
+        outputs = self.encoder(x, mask=mask#, src_key_padding_mask=x_key_padding_mask
+        ).permute(
             1, 0, 2
         )
 

seq2seq_time/models/transformer_autor.py

@@ -0,0 +1,73 @@
+from tqdm.auto import tqdm
+from torch import nn
+import torch
+from torch.nn import functional as F
+
+
+import fast_transformers
+from fast_transformers.builders import TransformerEncoderBuilder
+
+class TransformerAutoR(nn.Module):
+    def __init__(self, x_dim, y_dim, hidden_out_size=256, nlayers=8, n_heads=8, use_lstm=False, attention_dropout=0, dropout=0, min_std=0.01):
+        super().__init__()
+        self._min_std = min_std
+        self.use_lstm = use_lstm
+        hidden_out_size = hidden_out_size//n_heads
+
+        x_size = x_dim + y_dim
+        
+        # TODO embedd both X's the same
+        if use_lstm:            
+            self.x_emb = LSTMBlock(x_size, x_size)
+        
+        self.enc_emb = nn.Linear(x_size, hidden_out_size*n_heads)
+        self.encoder = fast_transformers.builders.TransformerEncoderBuilder.from_kwargs(
+            attention_type="causal-linear",
+            n_layers=nlayers,
+            n_heads=n_heads,
+            feed_forward_dimensions=hidden_out_size*8*n_heads,
+            query_dimensions=hidden_out_size,
+            value_dimensions=hidden_out_size,
+            attention_dropout=attention_dropout,
+            dropout=dropout,
+        ).get()
+        self.mean = nn.Linear(hidden_out_size*n_heads, y_dim)
+        self.std = nn.Linear(hidden_out_size*n_heads, y_dim)
+
+    def forward(self, past_x, past_y, future_x, future_y=None, mask_context=True, mask_target=True):
+        device = next(self.parameters()).device
+        B, S, _ = future_x.shape
+        future_y_fake = past_y[:, -1:, :].repeat(1, S, 1).to(device)
+        # future_y_fake = (
+        #     torch.ones(past_y.shape[0], future_x.shape[1], past_y.shape[2]).float().to(device) * 0
+        # )
+        context = torch.cat([past_x, past_y], -1)
+        target = torch.cat([future_x, future_y_fake], -1)
+        x = torch.cat([context, target * 1], 1).detach()        
+        
+        # LSTM
+        if self.use_lstm:  
+            x = self.x_emb(x)
+            # Size([B, T, Y]) -> Size([B, T, Y])
+        
+        # Embed
+        x = self.enc_emb(x)
+        
+        # requires  (B, C, hidden_dim)
+        steps = past_y.shape[1]
+        N = x.shape[1]
+        mask = fast_transformers.masking.TriangularCausalMask(N, device=device)
+        outputs = self.encoder(x, attn_mask=mask)[:, steps:, :]
+        
+        # Size([B, T, emb_dim])
+        mean = self.mean(outputs)
+        log_sigma = self.std(outputs)
+        sigma = self._min_std + (1 - self._min_std) * F.softplus(log_sigma)
+        y_dist = torch.distributions.Normal(mean, sigma)
+
+        return (
+            y_dist,
+            {}
+        )
+
+

seq2seq_time/models/transformer_process.py

@@ -31,8 +31,8 @@ class LatentEncoder(nn.Module):
         self.encoder = nn.TransformerEncoder(
             layer_enc, num_layers=num_layers, norm=encoder_norm
         )
-        self.mean = nn.Linear(hidden_size, latent_dim)
-        self.log_var = nn.Linear(hidden_size, latent_dim)
+        self.mean = nn.Linear(hidden_size*3, latent_dim)
+        self.log_var = nn.Linear(hidden_size*3, latent_dim)
         self._min_std = min_std
 
     def forward(self, x, y):
@@ -48,7 +48,13 @@ class LatentEncoder(nn.Module):
 
         r = self.encoder(x, mask=mask)
         r = r.permute(1, 0, 2)  # (S,B,hidden_size) -> (B,S,hidden_size)
-        r = r.mean(1) #  (B,S,hidden_size) ->  (B,hidden_size)
+
+        # Aggregation (max/mean/last)
+        r_mean = r.mean(1)  #  (B,S,hidden_size) ->  (B,hidden_size)
+        r_last = r[:, -1, :] 
+        r_max = r.max(1)[0]
+        r = torch.cat([r_mean, r_last, r_max], -1)
+        
         mean = self.mean(r)
         log_sigma = self.log_var(r)
         sigma = self._min_std + (1 - self._min_std) * torch.sigmoid(log_sigma * 0.5)
@@ -56,7 +62,6 @@ class LatentEncoder(nn.Module):
         return dist
 
 
-
 class Decoder(nn.Module):
     def __init__(
         self,

seq2seq_time/models/transformer_seq.py

@@ -2,6 +2,7 @@ import torch
 from torch import nn
 from torch.nn import functional as F
 
+from ..util import mask_upper_triangular
 
 class TransformerSeq(nn.Module):
     """

seq2seq_time/models/transformer_seq2seq.py

@@ -2,7 +2,7 @@ import torch
 from torch import nn
 from torch.nn import functional as F
 
-
+from ..util import mask_upper_triangular
 
 class TransformerSeq2Seq(nn.Module):
     def __init__(self, x_size, y_size, hidden_size=16, nhead=8, nlayers=2, attention_dropout=0, min_std=0.01, nan_value=0):
@@ -16,7 +16,7 @@ class TransformerSeq2Seq(nn.Module):
         encoder_norm = nn.LayerNorm(hidden_size)
         layer_enc = nn.TransformerEncoderLayer(
             d_model=hidden_size,
-            dim_feedforward=hidden_size*4,
+            dim_feedforward=hidden_size*8,
             dropout=attention_dropout,
             nhead=nhead,
             # activation
@@ -27,7 +27,7 @@ class TransformerSeq2Seq(nn.Module):
         
         layer_dec = nn.TransformerDecoderLayer(
             d_model=hidden_size,
-            dim_feedforward=hidden_size*4,
+            dim_feedforward=hidden_size*8,
             dropout=attention_dropout,
             nhead=nhead,
         )
@@ -67,7 +67,6 @@ class TransformerSeq2Seq(nn.Module):
         # In transformers the memory and future_x need to be the same length. Lets use a permutation invariant agg on the context
         # Then expand it, so it's available as we decode, conditional on future_x
         memory = memory.max(dim=0, keepdim=True)[0].expand_as(future_x)
-
         outputs = self.decoder(future_x, memory, tgt_key_padding_mask=tgt_key_padding_mask)
         
         # [T, B, emb_dim] -> [B, T, emb_dim]