Example 08: Saving and Loading¶

After training a model, you need to save it -- either for deployment (run inference without retraining) or to resume training later. TSFast provides three approaches at different levels of convenience:

1. save_model / load_model -- save the model for inference
2. save / load -- save the entire Learner to resume training
3. save_checkpoint / load_checkpoint -- fallback when the Learner contains unpicklable components (e.g. lambda loss functions)

Prerequisites¶

This notebook builds on Example 00 (training basics). Make sure TSFast is installed:

uv sync --extra dev

Setup¶

In [1]:

import tempfile
from pathlib import Path

import torch
from tsfast.inference import load_model
from tsfast.tsdata.benchmark import create_dls_silverbox
from tsfast.training import Learner, RNNLearner, fun_rmse

import tempfile from pathlib import Path import torch from tsfast.inference import load_model from tsfast.tsdata.benchmark import create_dls_silverbox from tsfast.training import Learner, RNNLearner, fun_rmse

In [2]:

tmpdir = Path(tempfile.mkdtemp())

tmpdir = Path(tempfile.mkdtemp())

Train a Quick Model¶

We train a small LSTM on the Silverbox benchmark. This gives us a trained Learner to demonstrate all three save/load approaches.

In [3]:

dls = create_dls_silverbox(bs=16, win_sz=500, stp_sz=10)
lrn = RNNLearner(dls, rnn_type='lstm', hidden_size=40, metrics=[fun_rmse])
lrn.fit_flat_cos(n_epoch=3)

dls = create_dls_silverbox(bs=16, win_sz=500, stp_sz=10) lrn = RNNLearner(dls, rnn_type='lstm', hidden_size=40, metrics=[fun_rmse]) lrn.fit_flat_cos(n_epoch=3)

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

Epoch 1/3: 100%|██████████| 300/300 [00:00<00:00, 638.16it/s, train=0.0183 | valid=0.0049 | fun_rmse=0.0103]

Epoch 1/3: 100%|██████████| 300/300 [00:00<00:00, 637.64it/s, train=0.0183 | valid=0.0049 | fun_rmse=0.0103]

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

Epoch 2/3: 100%|██████████| 300/300 [00:00<00:00, 911.81it/s, train=0.0047 | valid=0.0044 | fun_rmse=0.0099]

Epoch 2/3: 100%|██████████| 300/300 [00:00<00:00, 910.64it/s, train=0.0047 | valid=0.0044 | fun_rmse=0.0099]

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

Epoch 3/3: 100%|██████████| 300/300 [00:00<00:00, 946.00it/s, train=0.0038 | valid=0.0031 | fun_rmse=0.0096]

Epoch 3/3: 100%|██████████| 300/300 [00:00<00:00, 944.88it/s, train=0.0038 | valid=0.0031 | fun_rmse=0.0096]

---

Part 1: Save/Load Model for Inference¶

Use save_model when you only need to run the trained model for inference. It saves the full model object (including normalization layers) via torch.save. Loading returns a standalone nn.Module -- no Learner needed.

In [4]:

lrn.save_model(tmpdir / 'model.pth')

lrn.save_model(tmpdir / 'model.pth')

load_model is a standalone function in tsfast.inference that returns the model in eval mode, ready for inference.

In [5]:

model = load_model(tmpdir / 'model.pth')
print(type(model))

model = load_model(tmpdir / 'model.pth') print(type(model))

<class 'tsfast.models._core.scaling.ScaledModel'>

Run inference directly with the loaded model. The normalization layers are included, so you can pass raw data.

In [6]:

xb, yb = next(iter(dls.valid))
xb = xb.to(lrn.device)

with torch.no_grad():
    pred = model(xb)
    if isinstance(pred, tuple):
        pred = pred[0]

print(f"Input:      {xb.shape}")
print(f"Prediction: {pred.shape}")

xb, yb = next(iter(dls.valid)) xb = xb.to(lrn.device) with torch.no_grad(): pred = model(xb) if isinstance(pred, tuple): pred = pred[0] print(f"Input: {xb.shape}") print(f"Prediction: {pred.shape}")

Input:      torch.Size([16, 500, 1])
Prediction: torch.Size([16, 500, 1])

Verify that the loaded model produces the same predictions as the original.

In [7]:

preds_orig, _ = lrn.get_preds()

lrn_loaded = RNNLearner(dls, rnn_type='lstm', hidden_size=40)
lrn_loaded.model = model
preds_loaded, _ = lrn_loaded.get_preds()

max_diff = (preds_orig - preds_loaded).abs().max().item()
print(f"Max prediction difference: {max_diff:.2e}")
assert max_diff < 1e-5, "Predictions should match!"

preds_orig, _ = lrn.get_preds() lrn_loaded = RNNLearner(dls, rnn_type='lstm', hidden_size=40) lrn_loaded.model = model preds_loaded, _ = lrn_loaded.get_preds() max_diff = (preds_orig - preds_loaded).abs().max().item() print(f"Max prediction difference: {max_diff:.2e}") assert max_diff < 1e-5, "Predictions should match!"

Max prediction difference: 0.00e+00

---

Part 2: Save/Load Learner for Training Resume¶

Use save / load when you want to resume training later. This pickles the entire Learner state (model, optimizer, recorder, hyperparameters) except for dls, which you re-provide when loading.

In [8]:

lrn.save(tmpdir / 'learner.pth')

lrn.save(tmpdir / 'learner.pth')

load is a classmethod. The only argument besides the path is dls -- everything else (model architecture, optimizer state, training history) is restored from the saved file.

In [9]:

lrn2 = Learner.load(tmpdir / 'learner.pth', dls=dls)

lrn2 = Learner.load(tmpdir / 'learner.pth', dls=dls)

The recorder history is preserved -- our 3 training epochs are still there.

In [10]:

print(f"Recorder entries after load: {len(lrn2.recorder)}")
assert len(lrn2.recorder) == 3

print(f"Recorder entries after load: {len(lrn2.recorder)}") assert len(lrn2.recorder) == 3

Recorder entries after load: 3

Now resume training for 3 more epochs. The recorder continues from where we left off.

In [11]:

lrn2.fit_flat_cos(n_epoch=3)

lrn2.fit_flat_cos(n_epoch=3)

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

Epoch 1/3: 100%|██████████| 300/300 [00:00<00:00, 945.57it/s, train=0.0048 | valid=0.0050 | fun_rmse=0.0101]

Epoch 1/3: 100%|██████████| 300/300 [00:00<00:00, 944.17it/s, train=0.0048 | valid=0.0050 | fun_rmse=0.0101]

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

Epoch 2/3: 100%|██████████| 300/300 [00:00<00:00, 977.71it/s, train=0.0042 | valid=0.0048 | fun_rmse=0.0100]

Epoch 2/3: 100%|██████████| 300/300 [00:00<00:00, 976.47it/s, train=0.0042 | valid=0.0048 | fun_rmse=0.0100]

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

Epoch 3/3: 100%|██████████| 300/300 [00:00<00:00, 940.86it/s, train=0.0037 | valid=0.0030 | fun_rmse=0.0096]

Epoch 3/3: 100%|██████████| 300/300 [00:00<00:00, 939.58it/s, train=0.0037 | valid=0.0030 | fun_rmse=0.0096]

In [12]:

print(f"Total recorder entries: {len(lrn2.recorder)}")
assert len(lrn2.recorder) == 6  # 3 original + 3 resumed

print(f"Total recorder entries: {len(lrn2.recorder)}") assert len(lrn2.recorder) == 6 # 3 original + 3 resumed

Total recorder entries: 6

The validation loss should be lower after the additional training.

In [13]:

val_before = lrn2.recorder[2][1]  # last epoch before resume
val_after = lrn2.recorder[-1][1]  # last epoch after resume
print(f"Val loss before resume: {val_before:.4f}")
print(f"Val loss after resume:  {val_after:.4f}")

val_before = lrn2.recorder[2][1] # last epoch before resume val_after = lrn2.recorder[-1][1] # last epoch after resume print(f"Val loss before resume: {val_before:.4f}") print(f"Val loss after resume: {val_after:.4f}")

Val loss before resume: 0.0031
Val loss after resume:  0.0030

This approach also works for TbpttLearner -- the subclass type and all attributes (like sub_seq_len) are preserved automatically.

---

Part 3: Checkpoint Fallback¶

If save() fails because the Learner contains components that can't be pickled (e.g. lambda loss functions), use save_checkpoint / load_checkpoint instead. This saves only the model weights, optimizer state, and recorder -- you must reconstruct the Learner manually.

In [14]:

lrn.save_checkpoint(tmpdir / 'checkpoint.pth')

lrn.save_checkpoint(tmpdir / 'checkpoint.pth')

To load a checkpoint, first create a new Learner with the same architecture and configuration, then call load_checkpoint.

In [15]:

lrn3 = RNNLearner(dls, rnn_type='lstm', hidden_size=40, metrics=[fun_rmse])
lrn3.load_checkpoint(tmpdir / 'checkpoint.pth')

lrn3 = RNNLearner(dls, rnn_type='lstm', hidden_size=40, metrics=[fun_rmse]) lrn3.load_checkpoint(tmpdir / 'checkpoint.pth')

In [16]:

print(f"Recorder entries after checkpoint load: {len(lrn3.recorder)}")
assert len(lrn3.recorder) == 3

print(f"Recorder entries after checkpoint load: {len(lrn3.recorder)}") assert len(lrn3.recorder) == 3

Recorder entries after checkpoint load: 3

Verify the checkpoint restored the same model weights.

In [17]:

preds_ckpt, _ = lrn3.get_preds()
max_diff = (preds_orig - preds_ckpt).abs().max().item()
print(f"Max prediction difference: {max_diff:.2e}")
assert max_diff < 1e-5, "Checkpoint should restore identical weights!"

preds_ckpt, _ = lrn3.get_preds() max_diff = (preds_orig - preds_ckpt).abs().max().item() print(f"Max prediction difference: {max_diff:.2e}") assert max_diff < 1e-5, "Checkpoint should restore identical weights!"

Max prediction difference: 0.00e+00

Resume training from the checkpoint -- the optimizer state is restored when fit() is called.

In [18]:

lrn3.fit_flat_cos(n_epoch=3)
print(f"Total recorder entries: {len(lrn3.recorder)}")
assert len(lrn3.recorder) == 6

lrn3.fit_flat_cos(n_epoch=3) print(f"Total recorder entries: {len(lrn3.recorder)}") assert len(lrn3.recorder) == 6

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

Epoch 1/3: 100%|██████████| 300/300 [00:00<00:00, 848.26it/s, train=0.0047 | valid=0.0036 | fun_rmse=0.0097]

Epoch 1/3: 100%|██████████| 300/300 [00:00<00:00, 847.21it/s, train=0.0047 | valid=0.0036 | fun_rmse=0.0097]

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

Epoch 2/3: 100%|██████████| 300/300 [00:00<00:00, 894.76it/s, train=0.0046 | valid=0.0051 | fun_rmse=0.0103]

Epoch 2/3: 100%|██████████| 300/300 [00:00<00:00, 893.72it/s, train=0.0046 | valid=0.0051 | fun_rmse=0.0103]

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

Epoch 3/3: 100%|██████████| 300/300 [00:00<00:00, 850.79it/s, train=0.0038 | valid=0.0030 | fun_rmse=0.0096]

Epoch 3/3: 100%|██████████| 300/300 [00:00<00:00, 849.83it/s, train=0.0038 | valid=0.0030 | fun_rmse=0.0096]

Total recorder entries: 6

Key Takeaways¶

Method	What it saves	When to use
save_model / load_model	Full model object	Inference / deployment (from tsfast.inference import load_model)
save / load	Entire Learner (minus dls)	Resume training (default)
save_checkpoint / load_checkpoint	Weights + optimizer + recorder	Resume training with unpicklable components

• Start with save / load -- it's the simplest approach for resuming training.
• Use save_model when you only need the model for inference and don't need the training loop.
• Fall back to save_checkpoint only if save() raises a pickling error (e.g. lambda loss functions).