seq_len and other fixes

models/modules/encoders.py

@@ -44,10 +44,10 @@ class LinBnDropSN(nn.Sequential):
 
 
 class InceptionEncoder(nn.Module):
-    def __init__(self, c_in, c_out, *args, **kwargs):
+    def __init__(self, c_in, c_out, dropout, layers, layer_size, *args, **kwargs):
         super().__init__()
         self.net = CausalInceptionTimePlus(
-            c_in=c_in, c_out=c_out, custom_head=custom_head, *args, **kwargs
+            c_in=c_in, c_out=c_out, ks=[39, 19, 3], custom_head=custom_head, coord=True, fc_dropout=dropout, bn=True, depth=layers, nf=layer_size, *args, **kwargs
         )
         bn = kwargs.get("bn", True)
         fc_dropout = kwargs.get("fc_dropout", 0.15)
@@ -59,7 +59,7 @@ class InceptionEncoder(nn.Module):
         )
         self.head = nn.Sequential(
             # just to make sure we get a spectral norm final layer (after cat)
-            LinBnDropSN(c_out*2, c_out*2, bn=bn, p=fc_dropout),
+            LinBnDropSN(c_out*2, c_out, bn=bn, p=fc_dropout),
         )
 
     def forward(self, x):
@@ -258,15 +258,17 @@ class MLPEncoder(nn.Module):
         super().__init__()
         self.net = INR(
             in_feats=c_in,
+            out_feats=layer_size,
             scales=scales,
             n_fourier_feats=n_fourier_feats,
             layers=layers,
             layer_size=layer_size,
         )
+        self.head = nn.Linear(layer_size, c_out)
 
     def forward(self, x):
         """
         Takes in a sequence of shape (batch, sequence, features)
         and outputs a representation of shape (batch, features)
         """
-        return self.net(x)[:, -1]
+        return self.head(self.net(x)[:, -1])

models/modules/inr.py

@@ -31,7 +31,7 @@ class INRLayer(nn.Module):
 
 
 class INR(nn.Module):
-    def __init__(self, in_feats: int, layers: int, layer_size: int, n_fourier_feats: int, scales: float,
+    def __init__(self, in_feats: int, out_feats:int, layers: int, layer_size: int, n_fourier_feats: int, scales: float,
                  dropout: Optional[float] = 0.1):
         super().__init__()
         self.features = nn.Linear(in_feats, layer_size) if n_fourier_feats == 0 \
@@ -39,7 +39,7 @@ class INR(nn.Module):
         in_size = layer_size if n_fourier_feats == 0 \
             else n_fourier_feats
         layers = [INRLayer(in_size, layer_size, dropout=dropout)] + \
-                 [INRLayer(layer_size, layer_size, dropout=dropout) for _ in range(layers - 1)]
+                 [INRLayer(layer_size, out_feats, dropout=dropout) for _ in range(layers - 1)]
         self.layers = nn.Sequential(*layers)
 
     def forward(self, x: Tensor) -> Tensor:

models/modules/inrplus2.py

@@ -12,16 +12,16 @@ from torch import Tensor
 from models.modules.feature_transforms import GaussianFourierFeatureTransform
 
 from tsai.models.InceptionTimePlus import InceptionTimePlus
-from .causalinception import CausalInceptionTimePlus, CausalConv1d
+from .causalinception import CausalInceptionTimePlus, CausalConv1d, Conv
 
 def custom_head(head_nf, c_out, seq_len):
     return nn.Sequential(
-        CausalConv1d(head_nf, c_out, 1, bias=False)
-        
+        # CausalConv1d(head_nf, c_out, 1, bias=False, norm="Spectral")
+        Conv(head_nf, c_out, 1, bias=False, norm="Spectral"),
     )
 
 class INRPlus2(nn.Module):
-    def __init__(self, in_feats: int, layers: int, layer_size: int, n_fourier_feats: int, scales: float,
+    def __init__(self, in_feats: int, out_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.n_fourier_feats = n_fourier_feats
@@ -31,8 +31,8 @@ class INRPlus2(nn.Module):
         in_size = in_feats if n_fourier_feats == 0 \
             else n_fourier_feats+in_feats
         self.layers = CausalInceptionTimePlus(
-            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
+            in_size, out_feats, seq_len=None, nf=layer_size, depth=layers,
+                 flatten=False, concat_pool=False, fc_dropout=dropout, conv_dropout=dropout/4, 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)] + \
         #          [INRPlusLayer(layer_size, layer_size, dropout=dropout) for _ in range(layers - 1)]

models/modules/metareghead.py

@@ -7,9 +7,16 @@ from models.modules.regressors import RidgeRegressor
 from models.modules.inr import INR, INRLayer
 
 class SumHead(nn.Module):
-    def __init__(self, d, c_out=1, ):
+    def __init__(self, d, c_out=1, dropout=0):
         super().__init__()
-        self.l = nn.Linear(d, c_out) # init a random transform
+        # self.conv = nn.Sequential(
+        #     CausalConv1d(head_nf, c_out, 1, bias=False, norm="Spectral"),
+        # )
+        self.l = nn.Sequential(
+            INRLayer(d, d, dropout=dropout),
+            # INRLayer(d, d, dropout=dropout),
+            nn.Linear(d, c_out)
+        ) # nn.Linear(d, c_out) # init a random transform
 
     def forward(self, query, support, support_labels):
         return self.l(query)
@@ -27,7 +34,7 @@ class TransformerHead(nn.Module):
             INRLayer(c_out, hidden_dim, dropout=0),
             nn.Linear(hidden_dim, hidden_dim)
         )
-        self.l = nn.MultiheadAttention(embed_dim=d, num_heads=num_heads, batch_first=True, kdim=d, vdim=hidden_dim, add_bias_kv=True, bias=True)
+        self.l = nn.MultiheadAttention(embed_dim=d, num_heads=num_heads, batch_first=True, kdim=d, vdim=hidden_dim, add_bias_kv=True, bias=True, dropout=0)
         # after using attention let's decode it
         self.decoder = nn.Sequential(
             INRLayer(d, d, dropout=dropout),
@@ -40,12 +47,12 @@ class TransformerHead(nn.Module):
         returns the classification score on the query set.
 
         Parameters:
-        query:  a (tasks_per_batch, n_query, d) Tensor.
-        support:  a (tasks_per_batch, n_support, d) Tensor.
-        support_labels: a (tasks_per_batch, n_support) Tensor.
-        n_way: a scalar. Represents the number of classes in a few-shot classification task.
-        n_shot: a scalar. Represents the number of support examples given per class.
-        lambda_reg: a scalar. Represents the strength of L2 regularization.
+            query:  a (tasks_per_batch, n_query, d) Tensor.
+            support:  a (tasks_per_batch, n_support, d) Tensor.
+            support_labels: a (tasks_per_batch, n_support) Tensor.
+            n_way: a scalar. Represents the number of classes in a few-shot classification task.
+            n_shot: a scalar. Represents the number of support examples given per class.
+            lambda_reg: a scalar. Represents the strength of L2 regularization.
         Returns: a (tasks_per_batch, n_query, n_way) Tensor.
         """
         # should be (batch, seq, feature)
@@ -62,7 +69,7 @@ class RegressionHead(nn.Module):
             # the regular DeepTime one
             self.head = RidgeRegressor()
         elif ("None" in base_learner):
-            self.head = SumHead(d=d)
+            self.head = SumHead(d=d, dropout=dropout)
         elif ("Transformer" in base_learner):
             self.head = TransformerHead(d=d, dropout=dropout, num_heads=num_heads)
         else: