Example 05: Loss Functions and Metrics¶

The choice of loss function and evaluation metrics significantly affects training behavior and how you assess model quality. TSFast provides specialized losses and metrics for time series system identification:

• Loss modifiers that wrap any base loss to skip transients, slice windows, or normalize scales
• Evaluation metrics standard in the system identification literature

This example walks through each one and explains when to use it.

Prerequisites¶

This example builds on concepts from Examples 00-02. Make sure you have completed those first.

Setup¶

In [1]:

import torch
from torch import nn

from tsfast.tsdata.benchmark import create_dls_silverbox
from tsfast.training import (
    RNNLearner,
    fun_rmse, nrmse, nrmse_std, mean_vaf,
    weighted_mae, norm_loss, cut_loss, mse_nan,
)

import torch from torch import nn from tsfast.tsdata.benchmark import create_dls_silverbox from tsfast.training import ( RNNLearner, fun_rmse, nrmse, nrmse_std, mean_vaf, weighted_mae, norm_loss, cut_loss, mse_nan, )

Load Dataset¶

We use the Silverbox benchmark throughout this example. Each model trains for only 5 epochs to keep things fast -- the focus here is on the loss and metric behavior, not on achieving the best possible fit.

In [2]:

dls = create_dls_silverbox(bs=16, win_sz=500, stp_sz=10)

dls = create_dls_silverbox(bs=16, win_sz=500, stp_sz=10)

The Default: MAE Loss¶

TSFast uses nn.L1Loss() (Mean Absolute Error) as the default training loss. MAE is more robust to outliers and measurement spikes than MSE (Mean Squared Error). MSE heavily penalizes large errors, which can cause the model to overfit to noisy data points. For system identification, where measurement noise and occasional spikes are common, MAE provides a more stable training signal.

In [3]:

lrn_mae = RNNLearner(dls, rnn_type='lstm', loss_func=nn.L1Loss(), metrics=[fun_rmse])
lrn_mae.show_batch(max_n=2)

lrn_mae = RNNLearner(dls, rnn_type='lstm', loss_func=nn.L1Loss(), metrics=[fun_rmse]) lrn_mae.show_batch(max_n=2)

In [4]:

lrn_mae.fit_flat_cos(n_epoch=5, lr=3e-3)

lrn_mae.fit_flat_cos(n_epoch=5, lr=3e-3)

Epoch 1/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 1/5:  50%|█████     | 151/300 [00:00<00:00, 301.24it/s]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 301.24it/s, train=0.0124 | valid=0.0051 | fun_rmse=0.0105]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 344.98it/s, train=0.0124 | valid=0.0051 | fun_rmse=0.0105]

Epoch 2/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 2/5:  67%|██████▋   | 200/300 [00:00<00:00, 398.76it/s]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 398.76it/s, train=0.0057 | valid=0.0044 | fun_rmse=0.0101]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 401.97it/s, train=0.0057 | valid=0.0044 | fun_rmse=0.0101]

Epoch 3/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 3/5:  68%|██████▊   | 204/300 [00:00<00:00, 406.32it/s]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 406.32it/s, train=0.0050 | valid=0.0051 | fun_rmse=0.0103]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 406.74it/s, train=0.0050 | valid=0.0051 | fun_rmse=0.0103]

Epoch 4/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 4/5:  68%|██████▊   | 204/300 [00:00<00:00, 407.91it/s]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 407.91it/s, train=0.0051 | valid=0.0063 | fun_rmse=0.0110]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 407.16it/s, train=0.0051 | valid=0.0063 | fun_rmse=0.0110]

Epoch 5/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 5/5:  68%|██████▊   | 204/300 [00:00<00:00, 407.06it/s]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 407.06it/s, train=0.0036 | valid=0.0030 | fun_rmse=0.0097]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 405.03it/s, train=0.0036 | valid=0.0030 | fun_rmse=0.0097]

Evaluation Metrics¶

While you train with a loss function, you evaluate with metrics. TSFast provides several standard metrics for system identification:

• fun_rmse -- Root Mean Square Error. The standard reporting metric in many fields. Penalizes large errors more than MAE because of the squaring.

• nrmse -- RMSE normalized by the variance of each target variable. This allows fair comparison across outputs with different scales. A value of 0 means perfect prediction; a value of 1 means the model is no better than predicting the mean.

• nrmse_std -- RMSE normalized by the standard deviation of each target variable. Similar to nrmse but uses std instead of variance for the denominator.

• mean_vaf -- Variance Accounted For, expressed as a percentage. Measures what fraction of the target signal's variance is explained by the model. 100% means perfect prediction. This metric is widely used in the system identification literature.

In [5]:

lrn = RNNLearner(dls, rnn_type='lstm', metrics=[fun_rmse, nrmse, nrmse_std, mean_vaf])
lrn.fit_flat_cos(n_epoch=5, lr=3e-3)

lrn = RNNLearner(dls, rnn_type='lstm', metrics=[fun_rmse, nrmse, nrmse_std, mean_vaf]) lrn.fit_flat_cos(n_epoch=5, lr=3e-3)

Epoch 1/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 1/5:  67%|██████▋   | 202/300 [00:00<00:00, 403.65it/s]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 403.65it/s, train=0.0122 | valid=0.0055 | fun_rmse=0.0104 | mean_vaf=96.3704 | nrmse=0.1906 | nrmse_std=0.0445]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 402.68it/s, train=0.0122 | valid=0.0055 | fun_rmse=0.0104 | mean_vaf=96.3704 | nrmse=0.1906 | nrmse_std=0.0445]

Epoch 2/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 2/5:  68%|██████▊   | 203/300 [00:00<00:00, 405.33it/s]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 405.33it/s, train=0.0056 | valid=0.0048 | fun_rmse=0.0101 | mean_vaf=96.7073 | nrmse=0.1853 | nrmse_std=0.0432]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 402.40it/s, train=0.0056 | valid=0.0048 | fun_rmse=0.0101 | mean_vaf=96.7073 | nrmse=0.1853 | nrmse_std=0.0432]

Epoch 3/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 3/5:  69%|██████▉   | 207/300 [00:00<00:00, 413.37it/s]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 413.37it/s, train=0.0052 | valid=0.0054 | fun_rmse=0.0105 | mean_vaf=96.2906 | nrmse=0.1926 | nrmse_std=0.0449]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 409.81it/s, train=0.0052 | valid=0.0054 | fun_rmse=0.0105 | mean_vaf=96.2906 | nrmse=0.1926 | nrmse_std=0.0449]

Epoch 4/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 4/5:  68%|██████▊   | 204/300 [00:00<00:00, 406.23it/s]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 406.23it/s, train=0.0052 | valid=0.0061 | fun_rmse=0.0109 | mean_vaf=95.9837 | nrmse=0.2005 | nrmse_std=0.0467]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 407.60it/s, train=0.0052 | valid=0.0061 | fun_rmse=0.0109 | mean_vaf=95.9837 | nrmse=0.2005 | nrmse_std=0.0467]

Epoch 5/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 5/5:  68%|██████▊   | 204/300 [00:00<00:00, 406.84it/s]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 406.84it/s, train=0.0036 | valid=0.0029 | fun_rmse=0.0095 | mean_vaf=96.9149 | nrmse=0.1756 | nrmse_std=0.0410]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 405.32it/s, train=0.0036 | valid=0.0029 | fun_rmse=0.0095 | mean_vaf=96.9149 | nrmse=0.1756 | nrmse_std=0.0410]

In [6]:

lrn.show_results(max_n=2)

lrn.show_results(max_n=2)

n_skip: Ignoring Transient Warmup¶

RNNs start from a zero hidden state. During the first few timesteps, the hidden state is "warming up" and predictions are unreliable. Setting n_skip on the Learner discards the first N timesteps from both loss and metric computation. This prevents the optimizer from wasting effort on the unavoidable warmup transient.

• n_skip=50 -- skip the first 50 timesteps when computing the loss

In [7]:

lrn_skip = RNNLearner(dls, rnn_type='lstm', n_skip=50, metrics=[fun_rmse])
lrn_skip.fit_flat_cos(n_epoch=5, lr=3e-3)

lrn_skip = RNNLearner(dls, rnn_type='lstm', n_skip=50, metrics=[fun_rmse]) lrn_skip.fit_flat_cos(n_epoch=5, lr=3e-3)

Epoch 1/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 1/5:  67%|██████▋   | 202/300 [00:00<00:00, 402.32it/s]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 402.32it/s, train=0.0101 | valid=0.0038 | fun_rmse=0.0048]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 401.07it/s, train=0.0101 | valid=0.0038 | fun_rmse=0.0048]

Epoch 2/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 2/5:  68%|██████▊   | 204/300 [00:00<00:00, 406.07it/s]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 406.07it/s, train=0.0037 | valid=0.0043 | fun_rmse=0.0053]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 404.74it/s, train=0.0037 | valid=0.0043 | fun_rmse=0.0053]

Epoch 3/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 3/5:  67%|██████▋   | 201/300 [00:00<00:00, 400.60it/s]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 400.60it/s, train=0.0034 | valid=0.0064 | fun_rmse=0.0084]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 400.42it/s, train=0.0034 | valid=0.0064 | fun_rmse=0.0084]

Epoch 4/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 4/5:  68%|██████▊   | 203/300 [00:00<00:00, 405.49it/s]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 405.49it/s, train=0.0033 | valid=0.0036 | fun_rmse=0.0048]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 403.51it/s, train=0.0033 | valid=0.0036 | fun_rmse=0.0048]

Epoch 5/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 5/5:  68%|██████▊   | 204/300 [00:00<00:00, 407.98it/s]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 407.98it/s, train=0.0014 | valid=0.0008 | fun_rmse=0.0019]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 407.11it/s, train=0.0014 | valid=0.0008 | fun_rmse=0.0019]

cut_loss: Evaluating a Window¶

cut_loss slices the sequence to a specific range before computing the loss. This is useful when you only care about predictions in a particular part of the sequence.

• l_cut=50 -- trim 50 timesteps from the left (start of the sequence)
• r_cut=450 -- keep up to timestep 450 (trim from the right)

This evaluates only timesteps 50 through 450 of each 500-step window.

A related modifier, rand_seq_len_loss, truncates every sequence in the minibatch to an individually random length before computing the loss, which exposes the model to varying prediction horizons during training.

In [8]:

my_cut_loss = cut_loss(nn.L1Loss(), l_cut=50, r_cut=450)
lrn_cut = RNNLearner(dls, rnn_type='lstm', loss_func=my_cut_loss, metrics=[fun_rmse])
lrn_cut.fit_flat_cos(n_epoch=5, lr=3e-3)

my_cut_loss = cut_loss(nn.L1Loss(), l_cut=50, r_cut=450) lrn_cut = RNNLearner(dls, rnn_type='lstm', loss_func=my_cut_loss, metrics=[fun_rmse]) lrn_cut.fit_flat_cos(n_epoch=5, lr=3e-3)

Epoch 1/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 1/5:  68%|██████▊   | 203/300 [00:00<00:00, 404.43it/s]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 404.43it/s, train=0.0114 | valid=0.0049 | fun_rmse=0.0114]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 403.09it/s, train=0.0114 | valid=0.0049 | fun_rmse=0.0114]

Epoch 2/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 2/5:  68%|██████▊   | 205/300 [00:00<00:00, 408.89it/s]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 408.89it/s, train=0.0041 | valid=0.0044 | fun_rmse=0.0111]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 406.19it/s, train=0.0041 | valid=0.0044 | fun_rmse=0.0111]

Epoch 3/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 3/5:  69%|██████▊   | 206/300 [00:00<00:00, 411.00it/s]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 411.00it/s, train=0.0032 | valid=0.0055 | fun_rmse=0.0120]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 408.10it/s, train=0.0032 | valid=0.0055 | fun_rmse=0.0120]

Epoch 4/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 4/5:  66%|██████▌   | 198/300 [00:00<00:00, 395.82it/s]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 395.82it/s, train=0.0035 | valid=0.0036 | fun_rmse=0.0105]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 396.49it/s, train=0.0035 | valid=0.0036 | fun_rmse=0.0105]

Epoch 5/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 5/5:  68%|██████▊   | 203/300 [00:00<00:00, 404.15it/s]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 404.15it/s, train=0.0018 | valid=0.0009 | fun_rmse=0.0097]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 403.25it/s, train=0.0018 | valid=0.0009 | fun_rmse=0.0097]

norm_loss: Scale-Invariant Training¶

When your system has multiple outputs with very different magnitudes (e.g., position in meters and velocity in m/s), the loss is dominated by the largest-scale output. norm_loss normalizes both predictions and targets before computing the loss, so all outputs contribute equally regardless of their physical scale.

norm_loss takes the output normalization statistics from the DataLoaders (dls.norm_stats.y) and uses them to normalize both prediction and target tensors before passing them to the base loss function.

In [9]:

my_norm_loss = norm_loss(nn.L1Loss(), dls.norm_stats.y)
lrn_norm = RNNLearner(dls, rnn_type='lstm', loss_func=my_norm_loss, metrics=[fun_rmse])
lrn_norm.fit_flat_cos(n_epoch=5, lr=3e-3)

my_norm_loss = norm_loss(nn.L1Loss(), dls.norm_stats.y) lrn_norm = RNNLearner(dls, rnn_type='lstm', loss_func=my_norm_loss, metrics=[fun_rmse]) lrn_norm.fit_flat_cos(n_epoch=5, lr=3e-3)

Epoch 1/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 1/5:  68%|██████▊   | 204/300 [00:00<00:00, 407.29it/s]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 407.29it/s, train=0.1974 | valid=0.0727 | fun_rmse=0.0096]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 404.91it/s, train=0.1974 | valid=0.0727 | fun_rmse=0.0096]

Epoch 2/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 2/5:  68%|██████▊   | 203/300 [00:00<00:00, 404.41it/s]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 404.41it/s, train=0.1025 | valid=0.0949 | fun_rmse=0.0102]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 404.08it/s, train=0.1025 | valid=0.0949 | fun_rmse=0.0102]

Epoch 3/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 3/5:  68%|██████▊   | 205/300 [00:00<00:00, 408.21it/s]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 408.21it/s, train=0.0956 | valid=0.0957 | fun_rmse=0.0103]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 406.24it/s, train=0.0956 | valid=0.0957 | fun_rmse=0.0103]

Epoch 4/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 4/5:  67%|██████▋   | 201/300 [00:00<00:00, 401.83it/s]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 401.83it/s, train=0.0854 | valid=0.0779 | fun_rmse=0.0098]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 398.75it/s, train=0.0854 | valid=0.0779 | fun_rmse=0.0098]

Epoch 5/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 5/5:  68%|██████▊   | 205/300 [00:00<00:00, 409.04it/s]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 409.04it/s, train=0.0628 | valid=0.0528 | fun_rmse=0.0095]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 406.87it/s, train=0.0628 | valid=0.0528 | fun_rmse=0.0095]

Weighted MAE¶

weighted_mae applies log-spaced weights along the time axis, giving higher weight to earlier timesteps and lower weight to later ones. This is useful when early dynamics matter more than steady-state behavior, for example when modeling transient responses or step responses where the initial trajectory is most informative.

In [10]:

lrn_wmae = RNNLearner(dls, rnn_type='lstm', loss_func=weighted_mae, metrics=[fun_rmse])
lrn_wmae.fit_flat_cos(n_epoch=5, lr=3e-3)

lrn_wmae = RNNLearner(dls, rnn_type='lstm', loss_func=weighted_mae, metrics=[fun_rmse]) lrn_wmae.fit_flat_cos(n_epoch=5, lr=3e-3)

Epoch 1/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 1/5:  68%|██████▊   | 204/300 [00:00<00:00, 407.06it/s]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 407.06it/s, train=0.0163 | valid=0.0094 | fun_rmse=0.0112]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 405.78it/s, train=0.0163 | valid=0.0094 | fun_rmse=0.0112]

Epoch 2/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 2/5:  68%|██████▊   | 203/300 [00:00<00:00, 404.62it/s]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 404.62it/s, train=0.0083 | valid=0.0081 | fun_rmse=0.0103]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 404.71it/s, train=0.0083 | valid=0.0081 | fun_rmse=0.0103]

Epoch 3/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 3/5:  67%|██████▋   | 202/300 [00:00<00:00, 402.89it/s]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 402.89it/s, train=0.0078 | valid=0.0080 | fun_rmse=0.0102]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 401.42it/s, train=0.0078 | valid=0.0080 | fun_rmse=0.0102]

Epoch 4/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 4/5:  68%|██████▊   | 203/300 [00:00<00:00, 404.96it/s]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 404.96it/s, train=0.0076 | valid=0.0073 | fun_rmse=0.0098]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 402.70it/s, train=0.0076 | valid=0.0073 | fun_rmse=0.0098]

Epoch 5/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 5/5:  67%|██████▋   | 201/300 [00:00<00:00, 401.03it/s]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 401.03it/s, train=0.0064 | valid=0.0062 | fun_rmse=0.0096]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 400.24it/s, train=0.0064 | valid=0.0062 | fun_rmse=0.0096]

NaN-Safe Losses¶

Real measurement data often contains gaps: sensor dropouts, invalid samples, or variable-length records padded with NaN. A standard loss returns NaN as soon as a single target value is NaN, and training collapses. TSFast provides NaN-safe alternatives that mask out invalid timesteps:

• mse_nan -- MSE that drops every timestep whose target contains NaN
• ignore_nan -- decorator that adds the same masking to any loss function of your own
• nan_mean -- wraps a per-element loss into a masked mean with static tensor shapes, keeping it compatible with CUDA graphs (see Example 18)

In [11]:

targ = torch.randn(4, 100, 1)
pred = targ + 0.1 * torch.randn_like(targ)
targ[:, ::10] = float('nan')

print(f"mse_loss: {nn.functional.mse_loss(pred, targ).item()}")
print(f"mse_nan:  {mse_nan(pred, targ).item():.4f}")

targ = torch.randn(4, 100, 1) pred = targ + 0.1 * torch.randn_like(targ) targ[:, ::10] = float('nan') print(f"mse_loss: {nn.functional.mse_loss(pred, targ).item()}") print(f"mse_nan: {mse_nan(pred, targ).item():.4f}")

mse_loss: nan
mse_nan:  0.0113

The plain MSE is unusable once NaNs enter the target, while mse_nan returns the mean over the valid timesteps only. Pass it as loss_func to train directly on gappy measurement data.

Custom Learning Rate Schedules¶

All examples so far use fit_flat_cos, which keeps the learning rate flat then cosine-decays it to zero. For more control, call fit() directly with a scheduler_fn. tsfast.training exports two schedule functions:

• sched_ramp -- constant at start, linear ramp to end, then constant
• sched_lin_p -- linear decay from start to end, reaching end by position p

The flat-then-cosine schedule behind fit_flat_cos can be reused directly via from tsfast.training.schedulers import sched_flat_cos.

scheduler_fn is a factory (optimizer, total_steps) -> scheduler that creates a PyTorch LR scheduler. The schedule functions return a multiplier for a given position in [0, 1].

In [12]:

from torch.optim.lr_scheduler import LambdaLR
from tsfast.training import sched_ramp

lrn_sched = RNNLearner(dls, rnn_type='lstm', metrics=[fun_rmse])
lrn_sched.fit(
    n_epoch=5,
    lr=3e-3,
    scheduler_fn=lambda opt, steps: LambdaLR(
        opt, lambda s: sched_ramp(start=1.0, end=0.01, pos=s / steps)
    ),
)

from torch.optim.lr_scheduler import LambdaLR from tsfast.training import sched_ramp lrn_sched = RNNLearner(dls, rnn_type='lstm', metrics=[fun_rmse]) lrn_sched.fit( n_epoch=5, lr=3e-3, scheduler_fn=lambda opt, steps: LambdaLR( opt, lambda s: sched_ramp(start=1.0, end=0.01, pos=s / steps) ), )

Epoch 1/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 1/5:  69%|██████▉   | 207/300 [00:00<00:00, 412.12it/s]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 412.12it/s, train=0.0120 | valid=0.0077 | fun_rmse=0.0120]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 408.42it/s, train=0.0120 | valid=0.0077 | fun_rmse=0.0120]

Epoch 2/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 2/5:  69%|██████▉   | 207/300 [00:00<00:00, 413.66it/s]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 413.66it/s, train=0.0049 | valid=0.0034 | fun_rmse=0.0096]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 409.82it/s, train=0.0049 | valid=0.0034 | fun_rmse=0.0096]

Epoch 3/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 3/5:  68%|██████▊   | 204/300 [00:00<00:00, 407.57it/s]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 407.57it/s, train=0.0032 | valid=0.0029 | fun_rmse=0.0095]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 404.63it/s, train=0.0032 | valid=0.0029 | fun_rmse=0.0095]

Epoch 4/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 4/5:  68%|██████▊   | 203/300 [00:00<00:00, 405.23it/s]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 405.23it/s, train=0.0028 | valid=0.0029 | fun_rmse=0.0095]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 405.66it/s, train=0.0028 | valid=0.0029 | fun_rmse=0.0095]

Epoch 5/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 5/5:  69%|██████▊   | 206/300 [00:00<00:00, 410.51it/s]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 410.51it/s, train=0.0028 | valid=0.0029 | fun_rmse=0.0095]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 408.25it/s, train=0.0028 | valid=0.0029 | fun_rmse=0.0095]

Using PyTorch Built-in Schedulers¶

Any torch.optim.lr_scheduler works as a scheduler_fn. For example, OneCycleLR implements the 1cycle policy (warmup then cosine decay) and controls the learning rate entirely via max_lr -- the lr passed to fit() sets the optimizer baseline but OneCycleLR overrides it immediately.

In [13]:

from torch.optim.lr_scheduler import OneCycleLR

lrn_onecycle = RNNLearner(dls, rnn_type='lstm', metrics=[fun_rmse])
lrn_onecycle.fit(
    n_epoch=5,
    lr=3e-3,
    scheduler_fn=lambda opt, steps: OneCycleLR(opt, max_lr=3e-3, total_steps=steps),
)

from torch.optim.lr_scheduler import OneCycleLR lrn_onecycle = RNNLearner(dls, rnn_type='lstm', metrics=[fun_rmse]) lrn_onecycle.fit( n_epoch=5, lr=3e-3, scheduler_fn=lambda opt, steps: OneCycleLR(opt, max_lr=3e-3, total_steps=steps), )

Epoch 1/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 1/5:  68%|██████▊   | 203/300 [00:00<00:00, 405.83it/s]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 405.83it/s, train=0.0221 | valid=0.0053 | fun_rmse=0.0107]

Epoch 1/5: 100%|██████████| 300/300 [00:00<00:00, 404.56it/s, train=0.0221 | valid=0.0053 | fun_rmse=0.0107]

Epoch 2/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 2/5:  68%|██████▊   | 205/300 [00:00<00:00, 408.40it/s]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 408.40it/s, train=0.0080 | valid=0.0068 | fun_rmse=0.0113]

Epoch 2/5: 100%|██████████| 300/300 [00:00<00:00, 406.51it/s, train=0.0080 | valid=0.0068 | fun_rmse=0.0113]

Epoch 3/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 3/5:  69%|██████▊   | 206/300 [00:00<00:00, 411.82it/s]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 411.82it/s, train=0.0061 | valid=0.0039 | fun_rmse=0.0097]

Epoch 3/5: 100%|██████████| 300/300 [00:00<00:00, 408.45it/s, train=0.0061 | valid=0.0039 | fun_rmse=0.0097]

Epoch 4/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 4/5:  68%|██████▊   | 205/300 [00:00<00:00, 408.26it/s]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 408.26it/s, train=0.0036 | valid=0.0031 | fun_rmse=0.0096]

Epoch 4/5: 100%|██████████| 300/300 [00:00<00:00, 406.88it/s, train=0.0036 | valid=0.0031 | fun_rmse=0.0096]

Epoch 5/5:   0%|          | 0/300 [00:00<?, ?it/s]

Epoch 5/5:  68%|██████▊   | 204/300 [00:00<00:00, 407.36it/s]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 407.36it/s, train=0.0029 | valid=0.0030 | fun_rmse=0.0096]

Epoch 5/5: 100%|██████████| 300/300 [00:00<00:00, 405.63it/s, train=0.0029 | valid=0.0030 | fun_rmse=0.0096]

Key Takeaways¶

• MAE (default) is robust to outliers -- a good default for system identification where measurement noise and spikes are common.
• fun_rmse, nrmse, and mean_vaf are standard evaluation metrics. nrmse enables fair comparison across different-scale outputs, and mean_vaf reports the percentage of variance explained.
• n_skip excludes the RNN warmup transient from the loss, preventing the optimizer from fitting the unavoidable zero-state startup.
• cut_loss restricts the loss to a specific time window, useful when only part of the sequence matters.
• norm_loss enables scale-invariant training for multi-output systems by computing the loss in normalized space.
• weighted_mae emphasizes early timesteps, useful for transient-response modeling.
• mse_nan (with the ignore_nan and nan_mean wrappers) keeps training stable when targets contain NaNs from sensor dropouts or padding.
• fit() with scheduler_fn gives full control over the learning rate schedule. Use sched_ramp or sched_lin_p for custom warmup/decay profiles, or pass any PyTorch built-in scheduler like OneCycleLR.
• Choose metrics that match your evaluation requirements -- different applications call for different measures of model quality.