working

seq2seq_time/data/dataset.py

@@ -62,8 +62,8 @@ class Seq2SeqDataSet(torch.utils.data.Dataset):
         x_future = x[self.window_past:]
         y_future = y[self.window_past:]
 
-        # Stop it cheating by using future weather measurements
-        x_future[:, self._icol_blank] = 0
+        # Stop it cheating by using future weather measurements. Fill in with last value
+        x_future[:, self._icol_blank] = x_past[0, self._icol_blank]
         return x_past, y_past, x_future, y_future
 
 

seq2seq_time/models/baseline.py

@@ -0,0 +1,15 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+class BaselineLast(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.std = nn.Parameter(torch.tensor(1.))
+
+    def forward(self, past_x, past_y, future_x, future_y=None):
+        device = next(self.parameters()).device
+        B, S, F = future_x.shape
+        mean = past_y[:, -1:].repeat(1, S, 1)
+        std = (self.std * 1.0).repeat(1, S, 1)
+        return torch.distributions.Normal(mean, std)

seq2seq_time/models/lstm.py

@@ -0,0 +1,39 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+class LSTM(nn.Module):
+    def __init__(self, input_size, output_size, hidden_size=32, lstm_layers=2, lstm_dropout=0, _min_std = 0.05, nan_value=0):
+        super().__init__()
+        self._min_std = _min_std
+        self.nan_value = nan_value
+
+        self.lstm = nn.LSTM(
+            input_size=input_size + output_size,
+            hidden_size=hidden_size,
+            batch_first=True,
+            num_layers=lstm_layers,
+            dropout=lstm_dropout,
+        )
+        self.mean = nn.Linear(hidden_size, output_size)
+        self.std = nn.Linear(hidden_size, output_size)
+
+    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
+        )
+        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()
+        
+        steps = past_y.shape[1]
+        outputs, _ = self.lstm(x)
+        outputs = outputs[:, steps:, :]
+        
+        # outputs: [B, T, num_direction * H]
+        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/lstm_seq2seq.py

@@ -0,0 +1,39 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+class LSTMSeq2Seq(nn.Module):
+    def __init__(self, input_size, output_size, hidden_size=32, lstm_layers=2, lstm_dropout=0, _min_std = 0.05):
+        super().__init__()
+        self._min_std = _min_std
+
+        self.encoder = nn.LSTM(
+            input_size=input_size + output_size,
+            hidden_size=hidden_size,
+            batch_first=True,
+            num_layers=lstm_layers,
+            dropout=lstm_dropout,
+        )
+        self.decoder = nn.LSTM(
+            input_size=input_size,
+            hidden_size=hidden_size,
+            batch_first=True,
+            num_layers=lstm_layers,
+            dropout=lstm_dropout,
+        )
+        self.mean = nn.Linear(hidden_size, output_size)
+        self.std = nn.Linear(hidden_size, output_size)
+
+    def forward(self, past_x, past_y, future_x, future_y=None):
+        x = torch.cat([past_x, past_y], -1)
+        _, (h_out, cell) = self.encoder(x)
+        
+        # output = [batch size, seq len, hid dim * n directions]
+        outputs, (_, _) = self.decoder(future_x, (h_out, cell))
+        
+        # outputs: [B, T, num_direction * H]
+        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.py

@@ -0,0 +1,59 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+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):
+        super().__init__()
+        self._min_std = min_std
+        self.nan_value = nan_value
+        enc_x_dim = x_dim + y_dim
+
+        self.enc_emb = nn.Linear(enc_x_dim, hidden_size)
+        encoder_norm = nn.LayerNorm(hidden_size)
+        layer_enc = nn.TransformerEncoderLayer(
+            d_model=hidden_size,
+            dim_feedforward=hidden_size*4,
+            dropout=attention_dropout,
+            nhead=nhead,
+            # activation
+        )
+        self.encoder = nn.TransformerEncoder(
+            layer_enc, num_layers=nlayers, norm=encoder_norm
+        )
+        self.mean = nn.Linear(hidden_size, y_dim)
+        self.std = nn.Linear(hidden_size, y_dim)
+
+    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
+        )
+        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()
+
+        # Masks
+        x_mask = torch.isfinite(x) & (x != self.nan_value)
+        x[~x_mask] = 0
+        x = x.detach()
+        x_key_padding_mask = ~x_mask.any(-1)
+
+        x = self.enc_emb(x).permute(1, 0, 2)
+        
+        outputs = self.encoder(x, src_key_padding_mask=x_key_padding_mask).permute(
+            1, 0, 2
+        )
+
+        # Seems to help a little, especially with extrapolating out of bounds
+        steps = past_y.shape[1]
+        mean = self.mean(outputs)[:, steps:, :]
+        log_sigma = self.std(outputs)[:, steps:, :]
+        
+        sigma = self._min_std + (1 - self._min_std) * F.softplus(log_sigma)
+        return torch.distributions.Normal(mean, sigma)
+

seq2seq_time/models/transformer_seq2seq.py

@@ -0,0 +1,80 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+
+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):
+        super().__init__()
+        self._min_std = min_std
+        self.nan_value = nan_value
+        
+        self.enc_emb = nn.Linear(x_size + y_size, hidden_size)
+        self.dec_emb = nn.Linear(x_size, hidden_size)
+        
+        encoder_norm = nn.LayerNorm(hidden_size)
+        layer_enc = nn.TransformerEncoderLayer(
+            d_model=hidden_size,
+            dim_feedforward=hidden_size*4,
+            dropout=attention_dropout,
+            nhead=nhead,
+            # activation
+        )
+        self.encoder = nn.TransformerEncoder(
+            layer_enc, num_layers=nlayers, norm=encoder_norm
+        )
+        
+        layer_dec = nn.TransformerDecoderLayer(
+            d_model=hidden_size,
+            dim_feedforward=hidden_size*4,
+            dropout=attention_dropout,
+            nhead=nhead,
+        )
+        decoder_norm = nn.LayerNorm(hidden_size)
+        self.decoder = nn.TransformerDecoder(
+            layer_dec, num_layers=nlayers, norm=decoder_norm
+        )
+        self.mean = nn.Linear(hidden_size, y_size)
+        self.std = nn.Linear(hidden_size, y_size)
+
+
+    def forward(self, past_x, past_y, future_x, future_y=None):
+        device = next(self.parameters()).device
+        x = torch.cat([past_x, past_y], -1)
+
+        # Masks
+        future_mask = torch.isfinite(future_x) & (future_x!=self.nan_value)
+        tgt_key_padding_mask = ~future_mask.any(-1)
+
+        past_mask = torch.isfinite(x) & (x!=self.nan_value)
+        src_key_padding_mask = ~past_mask.any(-1)# * float('-inf')
+        
+        # Embed
+        x = self.enc_emb(x)
+        # Size([B, C, X]) -> Size([B, C, hidden_dim])
+        future_x = self.dec_emb(future_x)
+        # Size([B, C, T]) -> Size([B, C, hidden_dim])
+
+        x = x.permute(1, 0, 2)  # (B,C,hidden_dim) -> (C,B,hidden_dim)
+        future_x = future_x.permute(1, 0, 2) 
+        # requires  (C, B, hidden_dim)
+        memory = self.encoder(x, src_key_padding_mask=src_key_padding_mask)
+
+        # 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
+        # (C, B, emb_dim) -> (B, emb_dim) -> (T, B, emb_dim)
+        # 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]
+        outputs = outputs.permute(1, 0, 2).contiguous()
+        # 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)
+        return torch.distributions.Normal(mean, sigma)
+