LagLlamaEstimator
chronax.estimator.LagLlamaEstimator · inherits PyTorchLightningEstimator
An estimator training a ConvTSMixer model for forecasting. This class is uses the model defined in ConvTSMixerModel, and wraps it into a ConvTSMixerLightningModule for training purposes: training is performed using PyTorch Lightning's pl.Trainer class.
__init__(self, prediction_length, context_length=None, input_size=1, n_layer=1, n_embd_per_head=32, n_head=4, max_context_length=2048, rope_scaling=None, scaling='mean', lr=0.001, weight_decay=1e-08, aug_prob=0.1, freq_mask_rate=0.1, freq_mixing_rate=0.1, jitter_prob=0.0, jitter_sigma=0.03, scaling_prob=0.0, scaling_sigma=0.1, rotation_prob=0.0, permutation_prob=0.0, permutation_max_segments=5, permutation_seg_mode='equal', magnitude_warp_prob=0.0, magnitude_warp_sigma=0.2, magnitude_warp_knot=4, time_warp_prob=0.0, time_warp_sigma=0.2, time_warp_knot=4, window_slice_prob=0.0, window_slice_reduce_ratio=0.9, window_warp_prob=0.0, window_warp_window_ratio=0.1, window_warp_scales=[0.5, 2.0], distr_output='studentT', loss=NegativeLogLikelihood(), num_parallel_samples=100, batch_size=32, num_batches_per_epoch=50, trainer_kwargs=None, train_sampler=None, validation_sampler=None, time_feat=False, dropout=0.0, lags_seq=['QE', 'ME', 'W', 'D', 'h', 'min', 's'], data_id_to_name_map={}, use_cosine_annealing_lr=False, cosine_annealing_lr_args={}, track_loss_per_series=False, ckpt_path=None, nonnegative_pred_samples=False, use_single_pass_sampling=False, device=torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu'))
| Parameter | Type | Default | Description |
|---|---|---|---|
prediction_length |
int |
- | Length of the prediction horizon. |
context_length |
Optional[int] |
None |
Number of time steps prior to prediction time that the model takes as inputs (default: 10 * prediction_length). |
input_size |
int |
1 |
(undocumented) |
n_layer |
int |
1 |
(undocumented) |
n_embd_per_head |
int |
32 |
(undocumented) |
n_head |
int |
4 |
(undocumented) |
max_context_length |
int |
2048 |
(undocumented) |
rope_scaling |
- | None |
(undocumented) |
scaling |
Optional[str] |
"mean" |
(undocumented) |
lr |
float |
1e-3 |
Learning rate (default: 1e-3). |
weight_decay |
float |
1e-8 |
Weight decay regularization parameter (default: 1e-8). |
aug_prob |
float |
0.1 |
(undocumented) |
freq_mask_rate |
float |
0.1 |
(undocumented) |
freq_mixing_rate |
float |
0.1 |
(undocumented) |
jitter_prob |
float |
0.0 |
(undocumented) |
jitter_sigma |
float |
0.03 |
(undocumented) |
scaling_prob |
float |
0.0 |
(undocumented) |
scaling_sigma |
float |
0.1 |
(undocumented) |
rotation_prob |
float |
0.0 |
(undocumented) |
permutation_prob |
float |
0.0 |
(undocumented) |
permutation_max_segments |
int |
5 |
(undocumented) |
permutation_seg_mode |
str |
"equal" |
(undocumented) |
magnitude_warp_prob |
float |
0.0 |
(undocumented) |
magnitude_warp_sigma |
float |
0.2 |
(undocumented) |
magnitude_warp_knot |
int |
4 |
(undocumented) |
time_warp_prob |
float |
0.0 |
(undocumented) |
time_warp_sigma |
float |
0.2 |
(undocumented) |
time_warp_knot |
int |
4 |
(undocumented) |
window_slice_prob |
float |
0.0 |
(undocumented) |
window_slice_reduce_ratio |
float |
0.9 |
(undocumented) |
window_warp_prob |
float |
0.0 |
(undocumented) |
window_warp_window_ratio |
float |
0.1 |
(undocumented) |
window_warp_scales |
list |
[0.5, 2.0] |
(undocumented) |
distr_output |
str |
"studentT" |
Distribution to use to evaluate observations and sample predictions (default: StudentTOutput()). |
loss |
DistributionLoss |
NegativeLogLikelihood() |
Loss to be optimized during training (default: NegativeLogLikelihood()). |
num_parallel_samples |
int |
100 |
(undocumented) |
batch_size |
int |
32 |
The size of the batches to be used for training (default: 32). |
num_batches_per_epoch |
int |
50 |
Number of batches to be processed in each training epoch (default: 50). |
trainer_kwargs |
Optional[Dict[str, Any]] |
None |
Additional arguments to provide to pl.Trainer for construction. |
train_sampler |
Optional[InstanceSampler] |
None |
Controls the sampling of windows during training. |
validation_sampler |
Optional[InstanceSampler] |
None |
Controls the sampling of windows during validation. |
time_feat |
bool |
False |
(undocumented) |
dropout |
float |
0.0 |
(undocumented) |
lags_seq |
list |
["QE", "ME", "W", "D", "h", "min", "s"] |
(undocumented) |
data_id_to_name_map |
dict |
{} |
(undocumented) |
use_cosine_annealing_lr |
bool |
False |
(undocumented) |
cosine_annealing_lr_args |
dict |
{} |
(undocumented) |
track_loss_per_series |
bool |
False |
(undocumented) |
ckpt_path |
Optional[str] |
None |
(undocumented) |
nonnegative_pred_samples |
bool |
False |
(undocumented) |
use_single_pass_sampling |
bool |
False |
If True, use a single forward pass and sample N times from the saved distribution, much more efficient. If False, perform N forward passes and maintain N parallel prediction paths, this is true probalistic forecasting. (default: False) |
device |
torch.device |
torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") |
(undocumented) |
derive_auto_fields(cls, train_iter)
Derives automatic fields based on dataset statistics.
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
cls |
- | - | (undocumented) |
train_iter |
- | - | (undocumented) |
Returns: dict
create_transformation(self) -> Transformation
Creates the transformation pipeline for the dataset.
Parameters: None.
Returns: Transformation
create_lightning_module(self, use_kv_cache: bool = False) -> pl.LightningModule
Creates the PyTorch Lightning module used for training.
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
use_kv_cache |
bool |
False |
(undocumented) |
Returns: pl.LightningModule
create_training_data_loader(self, data: Dataset, module: LagLlamaLightningModule, shuffle_buffer_length: Optional[int] = None, **kwargs) -> Iterable
Creates the data loader for training.
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
data |
Dataset |
- | (undocumented) |
module |
LagLlamaLightningModule |
- | (undocumented) |
shuffle_buffer_length |
Optional[int] |
None |
(undocumented) |
**kwargs |
- | - | (undocumented) |
Returns: Iterable
create_validation_data_loader(self, data: Dataset, module: LagLlamaLightningModule, **kwargs) -> Iterable
Creates the data loader for validation.
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
data |
Dataset |
- | (undocumented) |
module |
LagLlamaLightningModule |
- | (undocumented) |
**kwargs |
- | - | (undocumented) |
Returns: Iterable
create_predictor(self, transformation: Transformation, module) -> PyTorchPredictor
Creates the predictor object used for inference.
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
transformation |
Transformation |
- | (undocumented) |
module |
- | - | (undocumented) |
Returns: PyTorchPredictor