partially converted to M2S (multi inputs)

models/DeepTIMe3.py

@@ -0,0 +1,88 @@
+ # Copyright (c) 2022, salesforce.com, inc.
+ # All rights reserved.
+ # SPDX-License-Identifier: BSD-3-Clause
+ # For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
+
+from typing import Optional
+
+import gin
+import torch
+import torch.nn as nn
+from torch import Tensor
+from einops import rearrange, repeat, reduce
+
+from models.modules.causalinception import CausalInceptionTimePlus, CausalConv1d
+from models.modules.inrplus2 import INRPlus2
+from models.modules.regressors import RidgeRegressor
+
+
+@gin.configurable()
+def deeptime3(dim_size:int, datetime_feats: int, layer_size: int, inr_layers: int, n_fourier_feats: int, scales: float):
+    return DeepTIMe3(dim_size, datetime_feats, layer_size, inr_layers, n_fourier_feats, scales)
+
+
+class DeepTIMe3(nn.Module):
+    def __init__(self, dim_size: int, datetime_feats: int, layer_size: int, inr_layers: int, n_fourier_feats: int, scales: float, dropout: float=0.3):
+        super().__init__()
+        
+        # encode the past
+        encoded_size = layer_size//2
+        self.encoder = CausalInceptionTimePlus(
+            c_in=dim_size, c_out=encoded_size, 
+            # nf=32, depth=6,
+            nf=17, depth=3,            
+            bn=True,
+            ks=[39, 19, 3],
+            coord=True, fc_dropout=dropout,
+        )
+        
+        # translate coords to a representation, given a summary of the past
+        coord_size = 1
+        in_feats=datetime_feats+encoded_size+coord_size
+        self.inr = INRPlus2(in_feats=in_feats, layers=inr_layers, layer_size=layer_size,
+                       n_fourier_feats=n_fourier_feats, scales=scales, dropout=dropout)
+        
+        # meta learn y given a representation
+        self.adaptive_weights = RidgeRegressor()
+
+        self.datetime_feats = datetime_feats
+        self.inr_layers = inr_layers
+        self.layer_size = layer_size
+        self.n_fourier_feats = n_fourier_feats
+        self.scales = scales
+        
+    def encode_and_decode(self, past_x, time, offset=0):
+        """
+        h_past = encode(past) # get representation of past
+        representation = decode(h_past, coords)
+        i = length of past, so we can offset the coords
+        """
+        encoded_x = self.encoder(past_x.transpose(2, 1))
+
+        # relative coordinates are the same for each batch, so we make them once and repeat them   
+        past_len = time.shape[1]
+        encoded_x = repeat(encoded_x, "b f -> b t f", t=past_len)
+        coords = self.get_coords(past_len).to(time.device) + offset
+        coords = repeat(coords, "1 t 1 -> b t 1", b=time.shape[0])
+        
+        
+        context_input = torch.cat([encoded_x, coords, time], dim=-1)
+
+        context_repr = self.inr(context_input)
+        return context_repr
+
+    def forward(self, context_past_x, context_y, query_past_x, query_y, context_time, query_time) -> Tensor:
+
+        context_reprs = self.encode_and_decode(context_past_x, context_time)
+        query_reprs = self.encode_and_decode(query_past_x, query_time, offset=context_reprs.shape[1])
+
+        w, b = self.adaptive_weights(context_reprs, context_y)
+        preds = self.forecast(query_reprs, w, b)
+        return preds
+
+    def forecast(self, inp: Tensor, w: Tensor, b: Tensor) -> Tensor:
+        return torch.einsum('... d o, ... t d -> ... t o', [w, inp]) + b
+    
+    def get_coords(self, lookback_len: int) -> Tensor:
+        coords = torch.linspace(0, 1, lookback_len)
+        return rearrange(coords, 't -> 1 t 1')

models/__init__.py

@@ -4,13 +4,16 @@ import torch
 
 from .DeepTIMe import deeptime
 from .DeepTIMe2 import deeptime2
+from .DeepTIMe3 import deeptime3
 
 
 def get_model(model_type: str, **kwargs: Union[int, float]) -> torch.nn.Module:
     if model_type == 'deeptime':
-        model = deeptime(datetime_feats=kwargs['datetime_feats'])
+        model = deeptime(datetime_feats=kwargs['datetime_feats'], dim_size=kwargs['dim_size'])
     elif model_type=="deeptime2":
-        model = deeptime2(datetime_feats=kwargs['datetime_feats'])
+        model = deeptime2(datetime_feats=kwargs['datetime_feats'], dim_size=kwargs['dim_size'])
+    elif model_type=="deeptime3":
+        model = deeptime3(datetime_feats=kwargs['datetime_feats'], dim_size=kwargs['dim_size'])
     else:
         raise ValueError(f"Unknown model type {model_type}")
     return model

models/modules/inrplus2.py

@@ -24,13 +24,14 @@ class INRPlus2(nn.Module):
     def __init__(self, in_feats: int, layers: int, layer_size: int, n_fourier_feats: int, scales: float,
                  dropout: Optional[float] = 0.5, bn=False, *args, **kwargs):
         super().__init__()
-        self.features = nn.Linear(in_feats, layer_size) if n_fourier_feats == 0 \
+        self.n_fourier_feats = n_fourier_feats
+        self.features = nn.Linear(in_feats, in_feats) if n_fourier_feats == 0 \
             else GaussianFourierFeatureTransform(in_feats, n_fourier_feats, scales)
-        in_size = layer_size if n_fourier_feats == 0 \
+        
+        in_size = in_feats if n_fourier_feats == 0 \
             else n_fourier_feats+in_feats
-        # import pdb; pdb.set_trace()
         self.layers = CausalInceptionTimePlus(
-            in_size-1, layer_size, seq_len=None, nf=layer_size, depth=layers,
+            in_size, layer_size, seq_len=None, nf=layer_size, depth=layers,
                  flatten=False, concat_pool=False, fc_dropout=dropout, conv_dropout=0.05, bn=bn, y_range=None, custom_head=custom_head, ks=[139, 19, 3], dilation=2, *args, **kwargs
         )
         # layers = [INRPlusLayer(in_size, layer_size, dropout=dropout)] + \
@@ -38,6 +39,7 @@ class INRPlus2(nn.Module):
         # self.layers = nn.Sequential(*layers)
 
     def forward(self, x: Tensor) -> Tensor:
-        x = self.features(x)
-        # import pdb; pdb.set_trace()
-        return self.layers(x.permute((0, 2, 1))).permute((0, 2, 1))
+        f = self.features(x)
+        if self.n_fourier_feats>0:
+            f = torch.concat([f, x], -1)
+        return self.layers(f.permute((0, 2, 1))).permute((0, 2, 1))