TrainState

train.TrainState

Lightweight training-state container (params + optimizer state).

__init__(self, params: dict, opt_state: optax.OptState, optimizer: optax.GradientTransformation)

Initializes the TrainState.

Parameter	Type	Default	Description
params	dict	-	(undocumented)
opt_state	optax.OptState	-	(undocumented)
optimizer	optax.GradientTransformation	-	(undocumented)

apply_gradients(self, grads: dict) -> TrainState

Applies gradients and updates the internal parameters and optimizer state.

| Parameter | Type | Default | Description | |-----------|---------------|-------------| | grads | dict | - | (undocumented) |

Returns: TrainState (the new training state).

make_loss_fn

train.make_loss_fn

Return a loss function (params, batch, rng) -> scalar.

Parameter	Type	Default	Description
model	nn.Module	-	(undocumented)

Returns: Callable

make_train_step

train.make_train_step

Return a JIT-compiled (state, batch, rng) -> (state, metrics).

Parameter	Type	Default	Description
model	nn.Module	-	(undocumented)
loss_fn	Callable	-	(undocumented)

Returns: Callable

train

train.train

Train a DeepAR model.

Parameter	Type	Default	Description
model	nn.Module	-	DeepAR_EncDec instance.
train_data	List	-	List of (y_series, exog_dict) tuples.
num_steps	int	-	Total training steps.
batch_size	int	-	Batch size.
input_size	int	100	History window length used to build batches.
h	int	24	Forecast horizon used to build batches.
learning_rate	float	1e-3	Peak learning rate.
weight_decay	float	1e-5	AdamW weight-decay.
warmup_steps	int	100	Linear warmup steps.
grad_clip	float	1.0	Gradient clipping norm.
valid_data	Optional[List]	None	Optional validation data.
early_stop_patience	int	-1	-1 disables early stopping.
seed	int	0	Random seed.
verbose	bool	True	Print progress.

Returns: Tuple[dict, List, List] (params, train_losses, valid_losses)