address #15

iTransformer/iTransformer.py

@@ -104,6 +104,7 @@ def __init__(
         ff_dropout = 0.,
         num_mem_tokens = 4,
         use_reversible_instance_norm = False,
+        reversible_instance_norm_affine = False,
         flash_attn = True
     ):
         super().__init__()
@@ -115,7 +116,7 @@ def __init__(
         pred_length = cast_tuple(pred_length)
         self.pred_length = pred_length
 
-        self.reversible_instance_norm = RevIN(num_variates) if use_reversible_instance_norm else None
+        self.reversible_instance_norm = RevIN(num_variates, affine = reversible_instance_norm_affine) if use_reversible_instance_norm else None
 
         self.layers = ModuleList([])
         for _ in range(depth):

iTransformer/iTransformer2D.py

@@ -160,6 +160,7 @@ def __init__(
         ff_dropout = 0.,
         num_mem_tokens = 4,
         use_reversible_instance_norm = False,
+        reversible_instance_norm_affine = True,
         flash_attn = True
     ):
         super().__init__()
@@ -175,7 +176,7 @@ def __init__(
         pred_length = cast_tuple(pred_length)
         self.pred_length = pred_length
 
-        self.reversible_instance_norm = RevIN(num_variates) if use_reversible_instance_norm else None
+        self.reversible_instance_norm = RevIN(num_variates, affine = reversible_instance_norm_affine) if use_reversible_instance_norm else None
 
         rotary_emb = RotaryEmbedding(dim_head)
 

iTransformer/revin.py

@@ -18,12 +18,17 @@
 # proposed in https://openreview.net/forum?id=cGDAkQo1C0p
 
 class RevIN(Module):
-    def __init__(self, num_variates, eps = 1e-5):
+    def __init__(
+        self,
+        num_variates,
+        affine = True,
+        eps = 1e-5
+    ):
         super().__init__()
         self.eps = eps
         self.num_variates = num_variates
-        self.gamma = nn.Parameter(torch.ones(num_variates, 1))
-        self.beta = nn.Parameter(torch.zeros(num_variates, 1))
+        self.gamma = nn.Parameter(torch.ones(num_variates, 1), requires_grad = affine)
+        self.beta = nn.Parameter(torch.zeros(num_variates, 1), requires_grad = affine)
 
     def forward(self, x, return_statistics = False):
         assert x.shape[1] == self.num_variates

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'iTransformer',
   packages = find_packages(exclude=[]),
-  version = '0.3.4',
+  version = '0.3.5',
   license='MIT',
   description = 'iTransformer - Inverted Transformer Are Effective for Time Series Forecasting',
   author = 'Phil Wang',