PyTorch module implementing the Chronos-2 model, ported from `amazon-science/chronos-forecasting `_ and adapted for Darts :class:`PLForecastingModule` interface. Parameters ---------- d_model
(
self,
d_model: int = 512,
d_kv: int = 64,
d_ff: int = 2048,
num_layers: int = 6,
num_heads: int = 8,
dropout_rate: float = 0.1,
layer_norm_epsilon: float = 1e-6,
feed_forward_proj: str = "relu",
rope_theta: float = 10000.0,
attn_implementation: Literal["eager", "sdpa"] | None = None,
chronos_config: dict[str, Any] | None = None,
**kwargs,
)
| 52 | |
| 53 | class _Chronos2Module(PLForecastingModule): |
| 54 | def __init__( |
| 55 | self, |
| 56 | d_model: int = 512, |
| 57 | d_kv: int = 64, |
| 58 | d_ff: int = 2048, |
| 59 | num_layers: int = 6, |
| 60 | num_heads: int = 8, |
| 61 | dropout_rate: float = 0.1, |
| 62 | layer_norm_epsilon: float = 1e-6, |
| 63 | feed_forward_proj: str = "relu", |
| 64 | rope_theta: float = 10000.0, |
| 65 | attn_implementation: Literal["eager", "sdpa"] | None = None, |
| 66 | chronos_config: dict[str, Any] | None = None, |
| 67 | **kwargs, |
| 68 | ): |
| 69 | """PyTorch module implementing the Chronos-2 model, ported from |
| 70 | `amazon-science/chronos-forecasting <https://github.com/amazon-science/chronos-forecasting>`_ and |
| 71 | adapted for Darts :class:`PLForecastingModule` interface. |
| 72 | |
| 73 | Parameters |
| 74 | ---------- |
| 75 | d_model |
| 76 | Dimension of the model embeddings, also called "model size" in Transformer. |
| 77 | d_kv |
| 78 | Dimension of the key and value projections in multi-head attention. |
| 79 | d_ff |
| 80 | Dimension of the feed-forward network hidden layer. |
| 81 | num_layers |
| 82 | Number of Chronos-2 encoder layers. |
| 83 | num_heads |
| 84 | Number of attention heads in each encoder block. |
| 85 | dropout_rate |
| 86 | Dropout rate of the model. |
| 87 | layer_norm_epsilon |
| 88 | Epsilon value for layer normalization layers. |
| 89 | feed_forward_proj |
| 90 | Activation of feed-forward network. |
| 91 | rope_theta |
| 92 | Base period for Rotary Position Embeddings (RoPE). |
| 93 | attn_implementation |
| 94 | Attention implementation to use. If None, defaults to "sdpa". |
| 95 | chronos_config |
| 96 | Configuration parameters for Chronos-2 model. See :class:`_Chronos2ForecastingConfig` for details. |
| 97 | **kwargs |
| 98 | all parameters required for :class:`darts.models.forecasting.pl_forecasting_module.PLForecastingModule` |
| 99 | base class. |
| 100 | """ |
| 101 | # for fine-tuning, model should be trained on pre-trained quantiles |
| 102 | enable_finetuning = kwargs.pop("enable_finetuning", False) |
| 103 | super().__init__(**kwargs) |
| 104 | self.d_model = d_model |
| 105 | self.d_kv = d_kv |
| 106 | self.d_ff = d_ff |
| 107 | self.num_layers = num_layers |
| 108 | self.num_heads = num_heads |
| 109 | self.dropout_rate = dropout_rate |
| 110 | self.layer_norm_epsilon = layer_norm_epsilon |
| 111 | self.feed_forward_proj = feed_forward_proj |
nothing calls this directly
no test coverage detected