layernorm key / values for seasonal, growth, level, after key value projection in classification wrapper

Author: lucidrains

etsformer_pytorch/etsformer_pytorch.py

@@ -388,7 +388,8 @@ def __init__(
         level_kernel_size = 3,
         growth_kernel_size = 3,
         seasonal_kernel_size = 3,
-        dropout = 0.
+        dropout = 0.,
+        norm_time_features_kv = False
     ):
         super().__init__()
         assert isinstance(etsformer, ETSFormer)
@@ -403,12 +404,14 @@ def __init__(
         self.queries = nn.Parameter(torch.randn(heads, dim_head))
 
         self.growth_to_kv = nn.Sequential(
+            nn.LayerNorm(model_dim),
             Rearrange('b n d -> b d n'),
             nn.Conv1d(model_dim, inner_dim * 2, growth_kernel_size, bias = False, padding = growth_kernel_size // 2),
             Rearrange('... (kv h d) n -> kv ... h n d', kv = 2, h = heads)
         )
 
         self.seasonal_to_kv = nn.Sequential(
+            nn.LayerNorm(model_dim),
             Rearrange('b n d -> b d n'),
             nn.Conv1d(model_dim, inner_dim * 2, seasonal_kernel_size, bias = False, padding = seasonal_kernel_size // 2),
             Rearrange('... (kv h d) n -> kv ... h n d', kv = 2, h = heads)
@@ -417,7 +420,8 @@ def __init__(
         self.level_to_kv = nn.Sequential(
             Rearrange('b n t -> b t n'),
             nn.Conv1d(time_features, inner_dim * 2, level_kernel_size, bias = False, padding = level_kernel_size // 2),
-            Rearrange('b (kv h d) n -> kv b h n d', kv = 2, h = heads)
+            Rearrange('b (kv h d) n -> kv b h n d', kv = 2, h = heads),
+            nn.LayerNorm(dim_head) if norm_time_features_kv else nn.Identity()
         )
 
         self.to_out = nn.Linear(inner_dim, model_dim)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'ETSformer-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.12',
+  version = '0.0.15',
   license='MIT',
   description = 'ETSTransformer - Exponential Smoothing Transformer for Time-Series Forecasting - Pytorch',
   author = 'Phil Wang',