return attention weights

models/etsformer/decoder.py

@@ -5,10 +5,11 @@ from einops import rearrange, reduce, repeat
 
 class DampingLayer(nn.Module):
 
-    def __init__(self, pred_len, nhead, dropout=0.1):
+    def __init__(self, pred_len, nhead, dropout=0.1, output_attention=False):
         super().__init__()
         self.pred_len = pred_len
         self.nhead = nhead
+        self.output_attention = output_attention
         self._damping_factor = nn.Parameter(torch.randn(1, nhead))
         self.dropout = nn.Dropout(dropout)
 
@@ -22,7 +23,10 @@ class DampingLayer(nn.Module):
         damping_factors = damping_factors.cumsum(dim=0)
         x = x.view(b, t, self.nhead, -1)
         x = self.dropout(x) * damping_factors.unsqueeze(-1)
-        return x.view(b, t, d)
+        x = x.view(b, t, d)
+        if self.output_attention:
+            return x, damping_factors
+        return x, None
 
     @property
     def damping_factor(self):
@@ -31,22 +35,26 @@ class DampingLayer(nn.Module):
 
 class DecoderLayer(nn.Module):
 
-    def __init__(self, d_model, nhead, c_out, pred_len, dropout=0.1):
+    def __init__(self, d_model, nhead, c_out, pred_len, dropout=0.1, output_attention=False):
         super().__init__()
         self.d_model = d_model
         self.nhead = nhead
         self.c_out = c_out
         self.pred_len = pred_len
+        self.output_attention = output_attention
 
-        self.growth_damping = DampingLayer(pred_len, nhead, dropout=dropout)
+        self.growth_damping = DampingLayer(pred_len, nhead, dropout=dropout, output_attention=output_attention)
         self.dropout1 = nn.Dropout(dropout)
 
     def forward(self, growth, season):
-        growth_horizon = self.growth_damping(growth[:, -1:])
+        growth_horizon, growth_damping = self.growth_damping(growth[:, -1:])
         growth_horizon = self.dropout1(growth_horizon)
 
         seasonal_horizon = season[:, -self.pred_len:]
-        return growth_horizon, seasonal_horizon
+
+        if self.output_attention:
+            return growth_horizon, seasonal_horizon, growth_damping
+        return growth_horizon, seasonal_horizon, None
 
 
 class Decoder(nn.Module):
@@ -64,11 +72,13 @@ class Decoder(nn.Module):
     def forward(self, growths, seasons):
         growth_repr = []
         season_repr = []
+        growth_dampings = []
 
         for idx, layer in enumerate(self.layers):
-            growth_horizon, season_horizon = layer(growths[idx], seasons[idx])
+            growth_horizon, season_horizon, growth_damping = layer(growths[idx], seasons[idx])
             growth_repr.append(growth_horizon)
             season_repr.append(season_horizon)
+            growth_dampings.append(growth_damping)
         growth_repr = sum(growth_repr)
         season_repr = sum(season_repr)
-        return self.pred(growth_repr), self.pred(season_repr)
+        return self.pred(growth_repr), self.pred(season_repr), growth_dampings