feat: 'README'

modeling_flame.py

@@ -550,8 +550,8 @@ class TimeDelayEmbedding(nn.Module):
         embedding = rearrange(embedding, '(b n) d -> b n d', b=b, n=n)
         return embedding
 
-    def forward(self, x, period=None, num_samples=1):
-        if not self.training and num_samples == 1:
+    def forward(self, x, period=None):
+        if not self.training:
             return self._embedding(x, period=period)
 
         if period is None:
@@ -604,14 +604,14 @@ class MixedEmbedding(nn.Module):
         proj.weight.data = resample(old=self.patch_embedding.proj.weight.data, new_patch_len=patch_len)
 
         patch_embedding = proj(patches)
-        time_delay_embedding = self.time_delay_embedding(x, period=patch_len, num_samples=1)
+        time_delay_embedding = self.time_delay_embedding(x, period=patch_len)
         embedding = patch_embedding + time_delay_embedding
         return embedding
 
-    def forward(self, x, inference_patch_len=48, num_samples=1):
+    def forward(self, x, inference_patch_len=48):
         # do patching
         # padding for the original stride
-        if not self.training and num_samples == 1:
+        if not self.training:
             return self._flex_embedding(x, inference_patch_len)
 
         seq_len = x.shape[-1]
@@ -627,7 +627,7 @@ class MixedEmbedding(nn.Module):
         patch_embedding = self.patch_embedding(patches)
 
         # [batch_size, patch_num, d_model]
-        time_delay_embedding = self.time_delay_embedding(x, num_samples=num_samples)
+        time_delay_embedding = self.time_delay_embedding(x)
 
         embedding = patch_embedding + time_delay_embedding
 
@@ -720,9 +720,9 @@ class FLAMEModel(nn.Module):
 
     def _predict(self, input, pred_len, inference_patch_len=48, num_samples=None):
         if num_samples is not None and num_samples > 1:
-            return self._prob_predict(input, pred_len, num_samples)
+            return self._prob_predict(input, pred_len, num_samples, inference_patch_len)
 
-        x_enc = self.embedding(input, inference_patch_len, num_samples=1)
+        x_enc = self.embedding(input, inference_patch_len)
 
         x_rec = self.encoder(x_enc)
 
@@ -741,13 +741,13 @@ class FLAMEModel(nn.Module):
         point_forecasts = rearrange(dec_out, 'b n p -> b (n p)')
         return point_forecasts[:, :pred_len]
 
-    def _prob_predict(self, input, pred_len, num_samples=None):
+    def _prob_predict(self, input, pred_len, num_samples=None, inference_patch_len=48):
 
-        x_enc = self.embedding(input, num_samples=num_samples)
+        x_enc = self.embedding(input, inference_patch_len=inference_patch_len)
 
         x_rec = self.encoder(x_enc)
 
-        predict_token_num = math.ceil(pred_len / self.patch_len)
+        predict_token_num = math.ceil(pred_len / inference_patch_len)
         weights = self._get_weights(predict_token_num).unsqueeze(0).unsqueeze(-1).to(input.device)
         last_token = x_rec[:, -1:, :]
         x_enc = weights * last_token.repeat(1, predict_token_num, 1)
@@ -757,6 +757,11 @@ class FLAMEModel(nn.Module):
         dist = self._prob_head(x_dec)
 
         samples = dist.sample((num_samples,))
+
+        weights = torch.eye(self.patch_len, device=x_dec.device)
+        resampled_weights = resample(old=weights, new_patch_len=inference_patch_len).T
+
+        samples = F.linear(samples, resampled_weights)
         samples = rearrange(samples, 's (b n) p -> b s (n p) ', n=predict_token_num)[:, :, :pred_len]
 
         prob_forecasts = samples