freeze reqs before I install vb2

models/DeepTIMe3.py

@@ -20,38 +20,38 @@ from models.modules.encoders import LSTMEncoder, TransformerEncoder2, Transforme
 # 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, dropout: float, base_learner: str, encoder:str, inr: str, seq_len: int):
-    return DeepTIMe3(dim_size, datetime_feats, layer_size, inr_layers, n_fourier_feats, scales, dropout, base_learner, encoder, inr, seq_len)
+def deeptime3(dim_size:int, datetime_feats: int, layer_size: int, inr_layers: int, n_fourier_feats: int, scales: float, dropout: float, lrn: str, enc:str, inr: str, seq_len: int):
+    return DeepTIMe3(dim_size, datetime_feats, layer_size, inr_layers, n_fourier_feats, scales, dropout, lrn, enc, inr, seq_len)
 
 
 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, base_learner:str='Ridge', encoder:str='inception', inr:str='INR', seq_len: int=46):
+    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, lrn:str='Ridge', enc:str='inception', inr:str='INR', seq_len: int=46):
         super().__init__()
         
         # encode the past
         encoded_size = layer_size
         encoder_features = 24
         encoder_layers = 3
-        if encoder == 'inception':
-            self.encoder = InceptionEncoder(
+        if enc == 'inception':
+            self.enc = InceptionEncoder(
                 c_in=dim_size, c_out=encoded_size, dilation=6,
                 layer_size=17, layers=encoder_layers, dropout=dropout,
             )
-        elif encoder == 'lstm':
-            self.encoder = LSTMEncoder(c_in=dim_size, c_out=encoded_size, dropout=dropout, layers=encoder_layers, layer_size=24)
-        elif encoder == 'lstm2':
-            self.encoder = LSTMEncoder2(c_in=dim_size, c_out=encoded_size, dropout=dropout, layers=encoder_layers, layer_size=32, seq_len=seq_len)
-        elif encoder == 'mlp':
-            self.encoder = MLPEncoder(c_in=dim_size, c_out=encoded_size, dropout=dropout, layers=encoder_layers, layer_size=256)
-        elif encoder == 'transformer':
-            self.encoder = TransformerEncoder(c_in=dim_size, c_out=encoded_size, dropout=dropout, layers=encoder_layers, layer_size=256, seq_len=seq_len)
-        elif encoder == 'transformer2':
-            self.encoder = TransformerEncoder2(c_in=dim_size, c_out=encoded_size, dropout=dropout, layers=encoder_layers, layer_size=256, seq_len=seq_len)
-        elif encoder == 'none':
-            self.encoder = None
+        elif enc == 'lstm':
+            self.enc = LSTMEncoder(c_in=dim_size, c_out=encoded_size, dropout=dropout, layers=encoder_layers, layer_size=24)
+        elif enc == 'lstm2':
+            self.enc = LSTMEncoder2(c_in=dim_size, c_out=encoded_size, dropout=dropout, layers=encoder_layers, layer_size=32, seq_len=seq_len)
+        elif enc == 'mlp':
+            self.enc = MLPEncoder(c_in=dim_size, c_out=encoded_size, dropout=dropout, layers=encoder_layers, layer_size=256)
+        elif enc == 'transformer':
+            self.enc = TransformerEncoder(c_in=dim_size, c_out=encoded_size, dropout=dropout, layers=encoder_layers, layer_size=256, seq_len=seq_len)
+        elif enc == 'transformer2':
+            self.enc = TransformerEncoder2(c_in=dim_size, c_out=encoded_size, dropout=dropout, layers=encoder_layers, layer_size=128, seq_len=seq_len)
+        elif enc == 'none':
+            self.enc = None
             encoded_size = 0
         else:
-            raise NotADirectoryError(encoder)
+            raise NotADirectoryError(enc)
         
         # translate coords to a representation, given a summary of the past
         coord_size = 1
@@ -66,7 +66,7 @@ class DeepTIMe3(nn.Module):
             raise NotImplementedError(inr)
         
         # meta learn y given a representation
-        self.regressionhead = RegressionHead(base_learner=base_learner, d=layer_size, dropout=dropout)
+        self.regressionhead = RegressionHead(lrn=lrn, d=layer_size, dropout=dropout)
 
         self.datetime_feats = datetime_feats
         self.inr_layers = inr_layers
@@ -83,8 +83,8 @@ class DeepTIMe3(nn.Module):
         
         # we summarize the past into a single hidden layer. Then repeat it for each coordinate
         past_len = time.shape[1]
-        if self.encoder is not None:
-            encoded_x = self.encoder(past_x)
+        if self.enc is not None:
+            encoded_x = self.enc(past_x)
             encoded_x = repeat(encoded_x, "b f -> b t f", t=past_len)
         
 
@@ -93,7 +93,7 @@ class DeepTIMe3(nn.Module):
         coords = repeat(coords, "1 t 1 -> b t 1", b=time.shape[0])
         
         # combine and run INR to decode the representation
-        if self.encoder is not None:
+        if self.enc is not None:
             context_input = torch.cat([encoded_x, coords, time], dim=-1)
         else:
             context_input = torch.cat([coords, time], dim=-1)

models/modules/causalinception.py

@@ -126,7 +126,8 @@ class CausalInceptionTimePlus(nn.Sequential):
         dilations = np.array([max(1, d*dilation) for d in range(depth)])
         d=np.array([dilations**i for i in range(3)]).T
         rf = ((ks-1)*d).sum(0)
-        print(f"receptive field {rf}={ks-1}*{d}")
+        # print(f"receptive field {rf}={ks-1}*{d}")
+        print(f"receptive field {rf}")
 
     def create_head(self, nf, c_out, seq_len, flatten=False, concat_pool=False, fc_dropout=0., bn=False, y_range=None):
         if flatten:

models/modules/encoders.py

@@ -133,14 +133,14 @@ class TransformerEncoder2(nn.Module):
         super().__init__()
         # d_model (82) must be divisible by n_heads (4)
         layer_size = layer_size // n_heads * n_heads
-        d_model = layer_size // 2 
+        d_model = layer_size // 4 
         self.net = TSPerceiver(
             c_in=c_in,
             c_out=c_out,
             seq_len=seq_len,
             
             # cat_szs=0, n_cont=0, 
-            n_latents=layer_size, d_latent=layer_size//4, 
+            n_latents=layer_size, d_latent=d_model, 
             # d_context=None, 
             self_per_cross_attn=1,
                 #  share_weights=True, cross_n_heads=1, d_head=None, 
@@ -229,7 +229,7 @@ class LSTMEncoder2(nn.Module):
             depth=layers,
             lstm_dropout=conv_dropout,
             fc_dropout=dropout,
-            pre_norm=False, use_token=True, use_pe=True, 
+            pre_norm=False, use_token=False, use_pe=False, 
             use_bn=False, 
         )
 

models/modules/metareghead.py

@@ -63,17 +63,17 @@ class TransformerHead(nn.Module):
 
 
 class RegressionHead(nn.Module):
-    def __init__(self, base_learner='Ridge', d=512, enable_scale=True, dropout=0.1, num_heads=16):
+    def __init__(self, lrn='Ridge', d=512, enable_scale=True, dropout=0.1, num_heads=16):
         super().__init__()
-        if ('Ridge' in base_learner):
+        if ('Ridge' in lrn):
             # the regular DeepTime one
             self.head = RidgeRegressor()
-        elif ("None" in base_learner):
+        elif ("None" in lrn):
             self.head = SumHead(d=d, dropout=dropout)
-        elif ("Transformer" in base_learner):
+        elif ("Transformer" in lrn):
             self.head = TransformerHead(d=d, dropout=dropout, num_heads=num_heads)
         else:
-            raise NotImplementedError(base_learner)
+            raise NotImplementedError(lrn)
         
         # Add a learnable scale
         self.enable_scale = enable_scale