diff --git a/benchmark/timeseries/gift_eval/README.md b/benchmark/timeseries/gift_eval/README.md index d49f394c0..b0117fad6 100644 --- a/benchmark/timeseries/gift_eval/README.md +++ b/benchmark/timeseries/gift_eval/README.md @@ -13,14 +13,18 @@ pip install --no-deps "git+https://github.com/SalesforceAIResearch/gift-eval.git ## Run +- **`TimesFM3`**: + + ```bash + python -m benchmark.timeseries.gift_eval.main --model timesfm3 + ``` + + `TimesFM3` follows official evaluation with symmetric averaging over two estimators, contexts truncated to 15,360 time steps and up to 32 variates per forward pass. + > [!NOTE] > Weights of `TimesFM3` are distributed under the [TimesFM Non-Commercial License v1.0](https://huggingface.co/google/timesfm-3.0-pytorch/blob/main/LICENSE). > Review the license before running the `TimesFM3` benchmark, which will download its weights noninteractively. -```bash -python -m benchmark.timeseries.gift_eval.main --model timesfm3 -``` - Results will be written to `benchmark/timeseries/gift_eval_out//`, and already completed configurations are skipped. ## Evaluate diff --git a/benchmark/timeseries/gift_eval/predictor.py b/benchmark/timeseries/gift_eval/predictor.py index 4266a5517..52d801bc9 100644 --- a/benchmark/timeseries/gift_eval/predictor.py +++ b/benchmark/timeseries/gift_eval/predictor.py @@ -36,6 +36,8 @@ def predict( ) -> Iterator[QuantileForecast]: for batch in batcher(dataset, batch_size=self.batch_size): + cpu_generator = torch.Generator().manual_seed(42) + ys = [torch.from_numpy(entry["target"]) for entry in batch] ys = [torch.atleast_2d(y) for y in ys] xs = [] @@ -67,6 +69,13 @@ def predict( for i, x_i in enumerate(xs): x[i, :, -x_i.size(-1) :] = x_i + if ( + self.config.max_variates is not None + and x.size(1) > self.config.max_variates - 1 + ): # Keep room for at least one target: + perm = torch.randperm(x.size(1), generator=cpu_generator) + x = x[:, perm[: self.config.max_variates - 1]] + y = y.transpose(-1, -2).to( self.device, memory_format=torch.contiguous_format, @@ -76,17 +85,30 @@ def predict( memory_format=torch.contiguous_format, ) - out = self.model( - x_context=x, - y_context=y, - x_query=torch.empty( - (x.size(0), self.prediction_length, 0), - device=self.device, - ), - num_estimators=self.config.num_estimators, - ) - pred = out.numerical - pred = pred.view(y.size(0), self.prediction_length, y.size(2), -1) + size = y.size(-1) + if self.config.max_variates is not None: + size = min(size, self.config.max_variates - x.size(-1)) + + preds = [] + for start in range(0, y.size(-1), size): + y_i = y[..., start : start + size] + if y_i.size(-1) < size: + y_i = torch.cat([y_i, y[..., : size - y_i.size(-1)]], -1) + + out = self.model( + x_context=x, + y_context=y_i, + x_query=torch.empty( + (x.size(0), self.prediction_length, 0), + device=self.device, + ), + num_estimators=self.config.num_estimators, + ) + pred = out.numerical + pred = pred.view(y.size(0), self.prediction_length, size, -1) + preds.append(pred) + + pred = torch.cat(preds, dim=-2)[..., : y.size(-1), :] pred = pred.permute(0, 3, 1, 2).cpu().numpy() # [B, Q, R, C] levels = out.numerical_columns[: pred.shape[1]] levels = [str(float(f"0.{q.split('q')[-1]}")) for q in levels] diff --git a/benchmark/timeseries/model.py b/benchmark/timeseries/model.py index feb9b85db..2689b66aa 100644 --- a/benchmark/timeseries/model.py +++ b/benchmark/timeseries/model.py @@ -15,6 +15,7 @@ class ModelConfig: factory: Callable[[torch.device], sdm.models.ICLModel] num_estimators: int = 1 max_context_size: int | None = None + max_variates: int | None = None MODEL_CONFIGS = { @@ -26,5 +27,6 @@ class ModelConfig: ), num_estimators=2, max_context_size=15_360, + max_variates=32, ), }