last tweak; give a little time overlap when constituting time tokens

lucidrains · Oct 29, 2023 · 243b6f4 · 243b6f4
1 parent 99fe3c4
commit 243b6f4
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/iTransformer/iTransformer2D.py b/iTransformer/iTransformer2D.py
@@ -202,9 +202,13 @@ def __init__(
             nn.LayerNorm(dim)
         )
 
+        time_kernel_size = lookback_len // num_time_tokens
+
         self.to_time_tokens = nn.Sequential(
-            Rearrange('b v (t n) -> b v t n', t = num_time_tokens),
-            nn.Linear(lookback_len // num_time_tokens, dim),
+            Rearrange('b v n -> (b v) 1 n'),
+            nn.ConstantPad1d((time_kernel_size, 0), value = 0.),
+            nn.Conv1d(1, dim, time_kernel_size * 2),
+            Rearrange('(b v) d t -> b v t d', v = num_variates),
             nn.LayerNorm(dim)
         )
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'iTransformer',
   packages = find_packages(exclude=[]),
-  version = '0.3.2',
+  version = '0.3.3',
   license='MIT',
   description = 'iTransformer - Inverted Transformer Are Effective for Time Series Forecasting',
   author = 'Phil Wang',