import json
import os
import struct
from collections import Counter
from dataclasses import asdict, dataclass
import torch
from torch import nn
from torch.nn import functional as F
from torchvision.models import resnet18, ResNet18_Weights
from safetensors.torch import load_file, save_file
from d2l import torch as d2l5.6 Saving, Loading, and Pretrained Weights
A trained network is two separate things kept in two separate places. The code is the class you wrote: its layers, its forward pass, the config that sized it. The model state is the tensors that training filled in: weights, biases, and running statistics. Optimizer moments, random-number streams, and the step belong to the wider training state. A weight file saves model state; a resumable checkpoint saves full training state. The code stays in your source repository, under version control, exactly like any other Python. To bring a model back you need both halves: run the code to rebuild an empty network, then pour the saved state into it.
This split explains most of what follows. It is why a checkpoint cannot resurrect a model on its own, why the config object from Section 5.1 belongs inside the checkpoint, and why the format that stores the state matters once you start sharing files with people who do not have your code.
import json
import struct
import warnings
from dataclasses import asdict, dataclass
import numpy as np
from d2l import tensorflow as d2l
import tensorflow as tf
from safetensors.tensorflow import load_file, save_fileimport json
import os
import struct
from dataclasses import asdict, dataclass
import jax
from jax import numpy as jnp
from flax import nnx
import optax
import orbax.checkpoint as ocp
from safetensors.flax import load_file, save_file
from d2l import jax as d2l
from d2l.nnx_resnet import ResNet50import json
import os
import struct
from collections import Counter
from dataclasses import asdict, dataclass
from mxnet import autograd, gluon, np, npx
from mxnet.gluon import nn
from safetensors.numpy import load_file, save_file
npx.set_np()5.6.1 State, Not Code
The state of a network is a dictionary from parameter names to tensors, the state_dict of Section 5.2. Before we save a whole model, the warm-up is that the same save/load calls work on any tensors, and on the lists and dicts that hold them.
The state of a network is a collection of named variables, the weights of Section 5.2. Before we save a whole model, the warm-up is that np.save and np.load work on any tensor, and, through NumPy’s pickle fallback, on the dicts that hold them.
The state of a network is a tree of typed, named variables introduced in Section 5.2. Before we save a whole model, the warm-up is that jnp.save and jnp.load work on any array, and, through NumPy’s pickle fallback, on the dicts that hold them.
The state of a network is a dictionary from parameter names to arrays, the collect_params dictionary of Section 5.2. Before we save a whole model, the warm-up is that npx.save and npx.savez work on any array and on named collections of them; npx.load hands a saved collection back as a dict.
x = torch.arange(4)
torch.save({'x': x, 'y': torch.zeros(4)}, 'tensors.pt')
torch.load('tensors.pt', weights_only=True){'x': tensor([0, 1, 2, 3]), 'y': tensor([0., 0., 0., 0.])}x = tf.range(4)
np.save('tensors.npy', {'x': x, 'y': tf.zeros(4)})
np.load('tensors.npy', allow_pickle=True).item(){'x': <tf.Tensor: shape=(4,), dtype=int32, numpy=array([0, 1, 2, 3], dtype=int32)>,
'y': <tf.Tensor: shape=(4,), dtype=float32, numpy=array([0., 0., 0., 0.], dtype=float32)>}x = jnp.arange(4)
jnp.save('tensors.npy', {'x': x, 'y': jnp.zeros(4)})
jnp.load('tensors.npy', allow_pickle=True).item(){'x': Array([0, 1, 2, 3], dtype=int32),
'y': Array([0., 0., 0., 0.], dtype=float32)}x = np.arange(4)
npx.savez('tensors-mx', x=x, y=np.zeros(4))
npx.load('tensors-mx'){'x': array([0., 1., 2., 3.]), 'y': array([0., 0., 0., 0.])}A model’s state_dict is one such dictionary, built for you. The keys are the dotted paths through the module tree (hidden.weight, output.bias); the values are the tensors, buffers included. Here is the tree for a small MLP.
A model’s state is one such collection, built for you: net.weights is the list of its variables, and each variable’s path names its place in the layer tree (mlp/hidden/kernel, mlp/output/bias; Keras calls a weight matrix a kernel). Keras invents layer names like dense_3 when you do not choose them, so we name the layers explicitly to keep the paths stable across instances. Here is the tree for a small MLP.
A model’s parameter state is one such structure. Its paths follow the module graph (hidden.kernel, output.bias; Flax calls a weight matrix a kernel). Here is the tree for a small MLP.
A model’s state is one such dictionary, built for you: collect_params maps the dotted paths through the block tree (hidden.weight, output.bias) to the parameters, the running statistics of any BatchNorm layers included. Here is the tree for a small MLP.
class MLP(nn.Module):
def __init__(self):
super().__init__()
self.hidden = nn.LazyLinear(256)
self.output = nn.LazyLinear(10)
def forward(self, x):
return self.output(F.relu(self.hidden(x)))
net = MLP()
X = torch.randn(2, 20)
Y = net(X)
{name: tuple(t.shape) for name, t in net.state_dict().items()}{'hidden.weight': (256, 20),
'hidden.bias': (256,),
'output.weight': (10, 256),
'output.bias': (10,)}class MLP(tf.keras.Model):
def __init__(self):
super().__init__(name='mlp')
self.hidden = tf.keras.layers.Dense(256, name='hidden')
self.out = tf.keras.layers.Dense(10, name='output')
def call(self, x):
return self.out(tf.nn.relu(self.hidden(x)))
net = MLP()
X = tf.random.normal((2, 20))
Y = net(X)
{v.path: tuple(v.shape) for v in net.weights}{'mlp/hidden/kernel': (20, 256),
'mlp/hidden/bias': (256,),
'mlp/output/kernel': (256, 10),
'mlp/output/bias': (10,)}class MLP(nnx.Module):
def __init__(self, rngs=None):
rngs = nnx.Rngs(0) if rngs is None else rngs
self.hidden = nnx.Linear(20, 256, rngs=rngs)
self.output = nnx.Linear(256, 10, rngs=rngs)
def __call__(self, x):
return self.output(nnx.relu(self.hidden(x)))
net = MLP()
X = jax.random.normal(d2l.get_key(), (2, 20))
Y = net(X)
params = nnx.state(net, nnx.Param)
[(path, tuple(value.shape)) for path, value in params.flat_state()][(('hidden', 'bias'), (256,)),
(('hidden', 'kernel'), (20, 256)),
(('output', 'bias'), (10,)),
(('output', 'kernel'), (256, 10))]class MLP(nn.Block):
def __init__(self):
super().__init__()
self.hidden = nn.Dense(256, activation='relu')
self.output = nn.Dense(10)
def forward(self, x):
return self.output(self.hidden(x))
net = MLP()
net.initialize()
X = np.random.uniform(size=(2, 20))
Y = net(X)
{name: p.shape for name, p in net.collect_params().items()}{'hidden.weight': (256, 20),
'hidden.bias': (256,),
'output.weight': (10, 256),
'output.bias': (10,)}Nothing in this dictionary knows it came from a class called MLP. The names and shapes refill a network with a compatible architecture and naming structure. Renaming an attribute or changing the nesting can break that match even when the computation is equivalent; the final section shows how to migrate such keys explicitly.
5.6.2 safetensors: the Interchange Format
torch.save writes its files with Python’s pickle, which does not store data so much as a program that reconstructs data. Unpickling runs that program. For a file you wrote and never let out of your control this is harmless. For a file you downloaded it is a remote-code-execution vector: loading the weights can run whatever the author’s pickle stream tells it to.
np.save writes NumPy’s .npy format. For a single tensor that is pure data: a small header with the dtype and shape, then the raw bytes. The warm-up dict is another matter. NumPy can only store a dict by falling back to Python’s pickle, which does not store data so much as a program that reconstructs data, and loading runs that program; that is what allow_pickle=True opted into. For a file you wrote and never let out of your control this is harmless. For a file you downloaded it is a remote-code-execution vector: any object in the stream can name a callable for the loader to invoke, which is why NumPy refuses pickled contents by default (allow_pickle=False).
jnp.save writes NumPy’s .npy format. For a single array that is pure data: a small header with the dtype and shape, then the raw bytes. The warm-up dict is another matter. NumPy can only store a dict by falling back to Python’s pickle, which does not store data so much as a program that reconstructs data, and loading runs that program; that is what allow_pickle=True opted into. For a file you wrote and never let out of your control this is harmless. For a file you downloaded it is a remote-code-execution vector: any object in the stream can name a callable for the loader to invoke, which is why NumPy refuses pickled contents by default (allow_pickle=False).
npx.savez writes MXNet’s own array format: the saver accepts nothing but MXNet arrays, and the file holds only their names, dtypes, shapes, and raw bytes. There is no pickle stream and nothing to execute on load, so the warm-up file was already safe. The catch is reach rather than safety: hardly anything outside MXNet reads such a file, which matters once you hand weights to someone who does not run an archived framework.
class Tripwire:
def __reduce__(self):
return (print, ('*** payload executed while loading ***',))
torch.save(Tripwire(), 'tripwire.pt')
_ = torch.load('tripwire.pt', weights_only=False) # the payload runs here*** payload executed while loading ***Since version 2.6, PyTorch defaults torch.load to weights_only=True, which refuses any pickle opcode that is not a plain tensor. The same file is now rejected instead of executed.
try:
torch.load('tripwire.pt', weights_only=True) # the default in this torch
except Exception as e:
print(type(e).__name__, str(e).splitlines()[0])UnpicklingError Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint. The allowlist behind weights_only=True is defense in depth, not a sandbox: it has itself had bypasses patched.
TensorFlow’s own model files sidestep pickle: neither a TF checkpoint nor the Keras .weights.h5/.keras formats execute it on load. They are not automatically safe to download, though. A .keras archive can carry serialized code (a Lambda layer’s function), which is why Keras refuses to deserialize such code unless you pass safe_mode=False, and why a model file from an untrusted source deserves the same caution as a pickle.
allow_pickle=False is a refusal, not a fix: it keeps the loader safe by declining to load the very files, dicts of named parameters, that model sharing needs.
safetensors removes the problem at the root by having no program to run. As Figure 5.6.1 lays out byte by byte, a safetensors file is an 8-byte little-endian integer giving the header length, then that many bytes of JSON naming each tensor’s dtype, shape, and byte range, then the raw tensor bytes back to back. There is no opcode stream to interpret, so the format itself carries no executable program. Loaders must still validate malformed headers, offsets, and allocation sizes. It is also framework-neutral and memory-mappable, which is why model hubs commonly prefer it. Save and reload the MLP’s state through it and confirm the round trip is exact.
Here the variable paths do the naming, so the flat mapping is a one-line comprehension over net.weights. Two quirks of the tensorflow binding to know: it converts through NumPy in both directions, so load_file hands back constant tensors that you assign into a model’s variables, and save_file overwrites the values of the dict you pass it during that conversion, so give it a throwaway copy.
One mismatch to bridge first: safetensors stores a flat mapping from names to tensors, while NNX state forms a nested pytree. This MLP has only string-keyed paths, so two small helpers join them with dots and split them on the way back. For an arbitrary pytree, preserve typed list indices and escape dots in keys, or restore values against a template instead of using this shortcut.
There is no safetensors.mxnet module; the archived framework never grew a native binding. That turns out to be the instructive case: because the format is just names, dtypes, and bytes, the safetensors.numpy binding plus a pair of dict comprehensions closes the gap, .asnumpy() on the way out and np.array on the way back. No framework-specific machinery is involved, which is the format’s framework-neutrality made visible.
save_file(net.state_dict(), 'mlp.safetensors')
clone = MLP()
clone(X) # materialize the lazy layers first
clone.load_state_dict(load_file('mlp.safetensors'))
clone.eval()
torch.equal(clone(X), Y)Truestate = {v.path: v for v in net.weights}
save_file(dict(state), 'mlp-tf.safetensors') # dict(): save_file clobbers its arg
restored = load_file('mlp-tf.safetensors')
clone = MLP()
clone(X) # build the variables first
for v in clone.weights:
v.assign(restored[v.path])
tf.reduce_all(clone(X) == Y)<tf.Tensor: shape=(), dtype=bool, numpy=True>def flatten(tree):
if hasattr(tree, 'flat_state'):
leaves = tree.flat_state()
else:
leaves = jax.tree_util.tree_flatten_with_path(tree)[0]
leaves = [(tuple(getattr(k, 'key', getattr(k, 'idx', k))
for k in path), value)
for path, value in leaves]
return {'.'.join(map(str, path)): value for path, value in leaves}
def unflatten(flat):
tree = {}
for name, value in flat.items():
*parents, leaf = name.split('.')
node = tree
for p in parents:
node = node.setdefault(p, {})
node[leaf] = value
return tree
save_file(flatten(params), 'mlp-jax.safetensors')
restored = unflatten(load_file('mlp-jax.safetensors'))
clone = MLP(nnx.Rngs(1))
clone_state = nnx.state(clone, nnx.Param)
nnx.replace_by_pure_dict(clone_state, restored)
nnx.update(clone, clone_state)
exact = jax.tree_util.tree_all(jax.tree_util.tree_map(
jnp.array_equal, restored, nnx.to_pure_dict(params)))
exact, jnp.array_equal(clone(X), Y)(True, Array(True, dtype=bool))state = {name: p.data().asnumpy() for name, p in net.collect_params().items()}
save_file(state, 'mlp-mx.safetensors')
restored = load_file('mlp-mx.safetensors')
clone = MLP()
clone.load_dict({name: np.array(v) for name, v in restored.items()})
(clone(X) == Y).all()array(True)Because the header is plain JSON at a known offset, you can read it without the library and see there is no magic to the format.
with open('mlp.safetensors', 'rb') as f:
n = struct.unpack('<Q', f.read(8))[0] # header length, little-endian
header = json.loads(f.read(n))
header['hidden.weight']{'dtype': 'F32', 'shape': [256, 20], 'data_offsets': [1024, 21504]}with open('mlp-tf.safetensors', 'rb') as f:
n = struct.unpack('<Q', f.read(8))[0] # header length, little-endian
header = json.loads(f.read(n))
header['mlp/hidden/kernel']{'dtype': 'F32', 'shape': [20, 256], 'data_offsets': [1024, 21504]}with open('mlp-jax.safetensors', 'rb') as f:
n = struct.unpack('<Q', f.read(8))[0] # header length, little-endian
header = json.loads(f.read(n))
header['hidden.kernel']{'dtype': 'F32', 'shape': [20, 256], 'data_offsets': [1024, 21504]}with open('mlp-mx.safetensors', 'rb') as f:
n = struct.unpack('<Q', f.read(8))[0] # header length, little-endian
header = json.loads(f.read(n))
header['hidden.weight']{'dtype': 'F32', 'shape': [256, 20], 'data_offsets': [1024, 21504]}torch.save keeps its place for your own scratch files and for the older code you will still meet. safetensors is what you use to hand a model to anyone else.
np.save and Keras’s own .weights.h5 keep their place for your own scratch files and checkpoints. safetensors is what you use to hand a model to anyone else.
jnp.save keeps its place for your own scratch arrays and quick experiments. safetensors is what you use to hand a model to anyone else.
save_parameters keeps its place for your own scratch files and checkpoints. safetensors is what you use to hand a model to anyone else, and for an archived framework it is also the exit route: current safetensors bindings outside MXNet can read the file you just wrote.
5.6.3 Checkpointing a Training Run
A checkpoint you can resume from holds more than weights. Resuming means picking up the optimizer where it stopped, and Adam’s state is the running first and second moments of the gradients from Section 5.2. Drop them and the optimizer restarts its momentum from zero, so the first steps after a resume no longer behave like a continuation. A full checkpoint therefore bundles the model state, the optimizer state, the RNG state (so data shuffling and dropout continue the same stream), the step counter, and the config that sizes the model when you rebuild it. Figure 5.6.2 pairs each of those five compartments with the exact thing it restores on resume.
Two details separate a checkpoint from a corrupted file. First, keep the contents to tensors and primitives so the file loads under weights_only=True; a dataclass config goes in as a plain dict via asdict. Second, write atomically: save to a temporary path and os.replace it into place, so a crash mid-write leaves the previous good checkpoint untouched rather than a half-written one.
In TensorFlow the bundle already has a name. tf.train.Checkpoint takes the objects to track as keyword arguments, model, optimizer, and a step counter, walks their variables, and saves and restores them as one unit, Adam’s moment estimates included. Because this native already covers the job, the tensorflow tab defines no helper of its own; the calls below are the idiom as you would write it in any project. Saves are numbered (run-tf-1, run-tf-2, …), so a crash cannot corrupt the previous completed number. The config travels in a JSON sidecar named for that number. A tracked tf.random.Generator carries the random stream; the process-global generator is not captured automatically.
NNX exposes the model and optimizer state as pytrees, including the optimizer’s step counter. Orbax, the JAX checkpointing library, saves and restores such trees whole: its StandardCheckpointer commits a temporary directory by rename. The demonstration overwrites a fixed path with force=True so the notebook is rerunnable; production runs should write numbered steps through CheckpointManager and retain the previous completed step. Because these natives already cover the job, the jax tab defines no helper of its own; the calls below are the idiom as you would write it in any project. The config rides along as one more branch of the saved tree, as does the PRNG key, which in JAX is explicit data rather than hidden global state.
In MXNet the bundle is three files rather than one. save_parameters covers the model; the optimizer state lives with the gluon.Trainer, whose save_states writes Adam’s moments (though not per-parameter attributes such as lr_mult, which come from your code on rebuild), and writes them with pickle, so a trainer-states file deserves the same your-own-disk-only caution as any pickle. The step counter and config travel in a JSON sidecar. One compartment of Figure 5.6.2 stays empty: MXNet has no API to snapshot its random-number generators, only npx.random.seed to restart them, so a resumed run reseeds rather than continuing the old stream. The helper atomically replaces each component, but three renames are not one transaction. Production code should write a numbered prefix and update one LATEST manifest only after all three files are complete.
def save_checkpoint(path, model, optimizer, step, cfg=None):
"""Atomically write a resumable training checkpoint."""
ckpt = {'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'step': step,
'cpu_rng': torch.get_rng_state()}
if cfg is not None:
ckpt['cfg'] = asdict(cfg)
if torch.cuda.is_available():
ckpt['cuda_rng'] = torch.cuda.get_rng_state_all()
tmp = path + '.tmp'
torch.save(ckpt, tmp)
os.replace(tmp, path)
def load_checkpoint(path, model, optimizer=None):
"""Restore the state written by save_checkpoint; return the raw dict."""
ckpt = torch.load(path, weights_only=True)
model.load_state_dict(ckpt['model'])
if optimizer is not None:
optimizer.load_state_dict(ckpt['optimizer'])
torch.set_rng_state(ckpt['cpu_rng'])
if torch.cuda.is_available() and 'cuda_rng' in ckpt:
torch.cuda.set_rng_state_all(ckpt['cuda_rng'])
return ckptdef save_checkpoint(prefix, model, trainer, step, cfg=None):
"""Replace each component of a resumable MXNet checkpoint."""
meta = {'step': step}
if cfg is not None:
meta['cfg'] = asdict(cfg)
model.save_parameters(prefix + '.params.tmp')
trainer.save_states(prefix + '.states.tmp')
with open(prefix + '.json.tmp', 'w') as f:
json.dump(meta, f)
for ext in ('.params', '.states', '.json'):
os.replace(prefix + ext + '.tmp', prefix + ext)
def load_checkpoint(prefix, model, trainer=None):
"""Restore the state written by save_checkpoint; return the metadata."""
model.load_parameters(prefix + '.params')
if trainer is not None:
trainer.load_states(prefix + '.states')
with open(prefix + '.json') as f:
return json.load(f)Train a tiny regressor for a hundred steps and checkpoint it. The Config dataclass is what a rebuild reads to size the model, so it travels with the weights.
@dataclass
class Config:
in_dim: int = 20
hidden: int = 64
lr: float = 0.05
def build(cfg):
return nn.Sequential(nn.Linear(cfg.in_dim, cfg.hidden), nn.ReLU(),
nn.Linear(cfg.hidden, 1))
torch.manual_seed(1)
cfg = Config()
data = torch.randn(256, cfg.in_dim)
target = data @ torch.randn(cfg.in_dim, 1) + 0.1 * torch.randn(256, 1)
loss = nn.MSELoss()
def step(model, opt):
opt.zero_grad()
l = loss(model(data), target)
l.backward()
opt.step()
return l.item()
net = build(cfg)
opt = torch.optim.Adam(net.parameters(), lr=cfg.lr)
for _ in range(100):
step(net, opt)
save_checkpoint('run.pt', net, opt, step=100, cfg=cfg)
round(loss(net(data), target).item(), 4)0.0016@dataclass
class Config:
in_dim: int = 20
hidden: int = 64
lr: float = 0.05
def build(cfg):
return tf.keras.Sequential([
tf.keras.layers.Dense(cfg.hidden, activation='relu'),
tf.keras.layers.Dense(1)])
tf.keras.utils.set_random_seed(1)
cfg = Config()
data = tf.random.normal((256, cfg.in_dim))
target = data @ tf.random.normal((cfg.in_dim, 1)) + 0.1 * tf.random.normal((256, 1))
loss = tf.keras.losses.MeanSquaredError()
def step(model, opt):
with tf.GradientTape() as tape:
l = loss(target, model(data))
opt.apply_gradients(zip(tape.gradient(l, model.trainable_variables),
model.trainable_variables))
return float(l)
net = build(cfg)
net(data[:1]) # build the variables
opt = tf.keras.optimizers.Adam(cfg.lr)
for _ in range(100):
step(net, opt)
rng = tf.random.Generator.from_seed(1)
ckpt = tf.train.Checkpoint(model=net, optimizer=opt, step=tf.Variable(100),
rng=rng)
path = ckpt.save('run-tf') # a numbered save: 'run-tf-1'
with open(path + '.cfg.json', 'w') as f:
json.dump(asdict(cfg), f)
path, round(float(loss(target, net(data))), 4)('run-tf-1', 0.0023)@dataclass
class Config:
in_dim: int = 20
hidden: int = 64
lr: float = 0.05
def build(cfg):
return nnx.Sequential(
nnx.Linear(cfg.in_dim, cfg.hidden, rngs=nnx.Rngs(0)), nnx.relu,
nnx.Linear(cfg.hidden, 1, rngs=nnx.Rngs(1)))
def fresh_state(cfg):
model = build(cfg)
optimizer = nnx.Optimizer(model, optax.adam(cfg.lr), wrt=nnx.Param)
return model, optimizer
cfg = Config()
model = build(cfg)
data = jax.random.normal(d2l.get_key(), (256, cfg.in_dim))
target = (data @ jax.random.normal(d2l.get_key(), (cfg.in_dim, 1))
+ 0.1 * jax.random.normal(d2l.get_key(), (256, 1)))
def loss(model):
return jnp.mean((model(data) - target) ** 2)
@nnx.jit
def step(model, optimizer):
l, grads = nnx.value_and_grad(loss)(model)
optimizer.update(model, grads)
return l
model, optimizer = fresh_state(cfg)
train_key = jax.random.key(1)
for _ in range(100):
l = step(model, optimizer)
ckptr = ocp.StandardCheckpointer()
ckptr.save(os.path.abspath('run-jax'), # orbax wants an absolute path
{'model': nnx.to_pure_dict(nnx.state(model)),
'optimizer': nnx.to_pure_dict(nnx.state(optimizer)),
'rng': train_key,
'cfg': asdict(cfg)}, force=True)
int(optimizer.step), round(float(loss(model)), 4)(100, 0.0021)@dataclass
class Config:
in_dim: int = 20
hidden: int = 64
lr: float = 0.05
def build(cfg):
net = nn.Sequential()
net.add(nn.Dense(cfg.hidden, activation='relu'), nn.Dense(1))
net.initialize()
return net
npx.random.seed(1)
cfg = Config()
data = np.random.normal(size=(256, cfg.in_dim))
target = data @ np.random.normal(size=(cfg.in_dim, 1)) \
+ 0.1 * np.random.normal(size=(256, 1))
def loss(model):
return float(((model(data) - target) ** 2).mean())
def step(model, trainer):
with autograd.record():
l = ((model(data) - target) ** 2).mean()
l.backward()
trainer.step(batch_size=1) # the loss is already a mean
return float(l)
net = build(cfg)
trainer = gluon.Trainer(net.collect_params(), 'adam',
{'learning_rate': cfg.lr})
for _ in range(100):
step(net, trainer)
save_checkpoint('run-mx', net, trainer, step=100, cfg=cfg)
round(loss(net), 4)0.0018The restore is exact. Corrupt every parameter, load the checkpoint back, and the loss returns to where it was.
Note the shape of the restore call: restore patches the tracked objects in place, matching each variable by its route through the object graph (model, then the layer, then its kernel) rather than by name, and it returns a status object. assert_consumed() on that status checks that every saved value found a variable and every variable found a value, turning a silent partial restore into a loud error.
Orbax fills a template with the saved arrays. We then update freshly built NNX objects from the restored pure dictionaries. Rebuilding first checks that the architecture still agrees with the checkpoint structure.
Note the shape of the restore: load_parameters and load_states patch the model and trainer in place, matching parameters by the same dotted names the file stores. There is no template to fill and no status object to check; what checking you want, you write yourself, as the final section of this page does with its key-set diff.
with torch.no_grad():
for p in net.parameters():
p.add_(1.0) # wreck the weights
before = loss(net(data), target).item()
load_checkpoint('run.pt', net, opt)
after = loss(net(data), target).item()
f'perturbed {before:.2f} -> restored {after:.4f}''perturbed 47291.83 -> restored 0.0016'for v in net.trainable_variables:
v.assign_add(tf.ones_like(v)) # wreck the weights
before = float(loss(target, net(data)))
ckpt.restore(path).assert_consumed()
after = float(loss(target, net(data)))
f'perturbed {before:.2f} -> restored {after:.4f}''perturbed 53413.12 -> restored 0.0023'model_state = nnx.state(model)
nnx.update(model, jax.tree.map(lambda p: p + 1.0, model_state))
before = float(loss(model))
template_model, template_optimizer = fresh_state(cfg)
template = {
'model': nnx.to_pure_dict(nnx.state(template_model)),
'optimizer': nnx.to_pure_dict(nnx.state(template_optimizer)),
'rng': jax.random.key(0),
'cfg': asdict(Config())}
ckpt = ckptr.restore(os.path.abspath('run-jax'), template)
cfg = Config(**ckpt['cfg'])
train_key = ckpt['rng']
restored_model = nnx.state(model)
nnx.replace_by_pure_dict(restored_model, ckpt['model'])
nnx.update(model, restored_model)
restored_optimizer = nnx.state(optimizer)
nnx.replace_by_pure_dict(restored_optimizer, ckpt['optimizer'])
nnx.update(optimizer, restored_optimizer)
after = float(loss(model))
f'perturbed {before:.2f} -> restored {after:.4f}''perturbed 50460.26 -> restored 0.0021'for p in net.collect_params().values():
p.set_data(p.data() + 1.0) # wreck the weights
before = loss(net)
load_checkpoint('run-mx', net, trainer)
after = loss(net)
f'perturbed {before:.2f} -> restored {after:.4f}''perturbed 49008.15 -> restored 0.0018'Now the reason the optimizer state is in the file. Resume the run two ways from the same checkpoint: once restoring the optimizer, once with a fresh one holding only the weights. The network is near its minimum, so the correct continuation barely moves. A fresh Adam, with its moment estimates reset and its bias correction starting over, takes an oversized first step and overshoots.
One Keras 3 tripwire sits in the resume path. A freshly constructed Adam owns no slot variables; the moment estimates are created only when the optimizer is built. Restore into a fresh optimizer without building it first and the saved moments have no variables to land in, so assert_consumed() fails with an unresolved optimizer object. The fix is one line before the restore: opt.build(model.trainable_variables).
net_full = build(cfg)
opt_full = torch.optim.Adam(net_full.parameters(), lr=cfg.lr)
load_checkpoint('run.pt', net_full, opt_full) # weights + optimizer
full = [round(step(net_full, opt_full), 4) for _ in range(5)]
net_fresh = build(cfg)
load_checkpoint('run.pt', net_fresh, optimizer=None) # weights only
opt_fresh = torch.optim.Adam(net_fresh.parameters(), lr=cfg.lr)
fresh = [round(step(net_fresh, opt_fresh), 4) for _ in range(5)]
print('full optimizer:', full)
print('fresh optimizer:', fresh)full optimizer: [0.0016, 0.0016, 0.0015, 0.0015, 0.0015]
fresh optimizer: [0.0016, 3.8324, 0.2922, 1.0421, 1.845]with open(path + '.cfg.json') as f:
cfg = Config(**json.load(f))
net_full = build(cfg)
net_full(data[:1])
opt_full = tf.keras.optimizers.Adam(cfg.lr)
opt_full.build(net_full.trainable_variables) # create Adam's slots first
rng_full = tf.random.Generator.from_seed(0)
tf.train.Checkpoint(model=net_full, optimizer=opt_full,
step=tf.Variable(0), rng=rng_full).restore(path).assert_consumed()
full = [round(step(net_full, opt_full), 4) for _ in range(5)]
net_fresh = build(cfg)
net_fresh(data[:1])
tf.train.Checkpoint(model=net_fresh).restore(path).expect_partial() # weights only
opt_fresh = tf.keras.optimizers.Adam(cfg.lr)
fresh = [round(step(net_fresh, opt_fresh), 4) for _ in range(5)]
print('full optimizer:', full)
print('fresh optimizer:', fresh)full optimizer: [0.0023, 0.0022, 0.0022, 0.0021, 0.002]
fresh optimizer: [0.0023, 2.2475, 0.259, 1.1466, 0.7407]full, full_opt = fresh_state(cfg)
full_state, full_opt_state = nnx.state(full), nnx.state(full_opt)
nnx.replace_by_pure_dict(full_state, ckpt['model'])
nnx.replace_by_pure_dict(full_opt_state, ckpt['optimizer'])
nnx.update(full, full_state)
nnx.update(full_opt, full_opt_state)
full_losses = []
for _ in range(5):
l = step(full, full_opt)
full_losses.append(round(float(l), 4))
fresh, fresh_opt = fresh_state(cfg)
fresh_state_ = nnx.state(fresh)
nnx.replace_by_pure_dict(fresh_state_, ckpt['model'])
nnx.update(fresh, fresh_state_)
fresh_losses = []
for _ in range(5):
l = step(fresh, fresh_opt)
fresh_losses.append(round(float(l), 4))
print('full optimizer:', full_losses)
print('fresh optimizer:', fresh_losses)full optimizer: [0.0021, 0.002, 0.002, 0.002, 0.002]
fresh optimizer: [0.0021, 3.0551, 0.2507, 0.8011, 1.4134]net_full = build(cfg)
trainer_full = gluon.Trainer(net_full.collect_params(), 'adam',
{'learning_rate': cfg.lr})
load_checkpoint('run-mx', net_full, trainer_full) # weights + optimizer
full = [round(step(net_full, trainer_full), 4) for _ in range(5)]
net_fresh = build(cfg)
load_checkpoint('run-mx', net_fresh, trainer=None) # weights only
trainer_fresh = gluon.Trainer(net_fresh.collect_params(), 'adam',
{'learning_rate': cfg.lr})
fresh = [round(step(net_fresh, trainer_fresh), 4) for _ in range(5)]
print('full optimizer:', full)
print('fresh optimizer:', fresh)full optimizer: [0.0018, 0.0018, 0.0017, 0.0017, 0.0017]
fresh optimizer: [0.0018, 3.6124, 0.1789, 1.1589, 1.8473]The full-state run keeps descending; the weights-only run spikes and has to claw its way back. That transient is the cost of forgetting the optimizer, and it is why “just the weights” is not a resumable checkpoint.
For models too large to hold in memory, checkpoints are split across several files with an index, and torch.load(..., mmap=True) pages tensors off disk on demand instead of copying the whole file up front. Combined with meta-device construction and load_state_dict(..., assign=True), this loads such a model without ever allocating its randomly-initialized weights; Chapter 13 returns to the machinery when models get that big.
For models too large to hold in memory, the format already cooperates: a TF checkpoint is an index file plus data shards rather than one monolith, and restores are lazy, so a variable created later is filled from the file at creation time instead of everything materializing up front. Multi-host tf.distribute jobs build on the same machinery; Chapter 13 returns to it when models get that big.
For models too large to hold in memory, orbax already works at scale: a checkpoint is a directory of per-array files rather than one monolith, a restore can target a device sharding so each accelerator materializes only its own pieces, and multi-host jobs save and restore in parallel; Chapter 13 returns to the machinery when models get that big.
For models too large to hold in memory, MXNet offers nothing comparable: a .params file loads whole, and the project was archived before sharded checkpoints and memory-mapped loading became routine. Models at that scale are where you leave the framework; Chapter 13 returns to the machinery the others provide.
5.6.4 Loading Weights You Did Not Train
The most common reason to load a state_dict is that someone else produced it. You take a network trained on a large dataset and adapt it: keep the learned feature extractor, replace the final layer for your own labels. The mechanics are state_dict manipulation. torchvision serves the weights through a weights= enum, which also downloads the matching parameters the first time.
The most common reason to load weights is that someone else produced them. You take a network trained on a large dataset and adapt it: keep the learned feature extractor, replace the final layer for your own labels. keras.applications is the built-in zoo; weights='imagenet' downloads the matching parameters the first time, and include_top=False drops the 1000-class head so you can attach your own.
The most common reason to load parameters is that someone else produced them. You take a network trained on a large dataset and adapt it: keep the learned feature extractor, replace the final layer for your own labels. The book’s NNX ResNet-50 loader downloads pinned Microsoft ImageNet safetensors from the Hugging Face Hub and maps them into an NNX model. We use that real artifact, then reuse the small MLP file to show the key-by-key surgery without burying the mechanism under a ResNet-sized state tree.
The most common reason to load parameters is that someone else produced them. You take a network trained on a large dataset and adapt it: keep the learned feature extractor, replace the final layer for your own labels. gluon.model_zoo.vision is the built-in zoo; pretrained=True downloads the matching parameters the first time. One caveat before you rely on it: the download comes from the archived project’s file hosting, which still worked when this notebook last ran but carries no promise of staying up, so re-verify it before building anything on top.
net = resnet18(weights=ResNet18_Weights.DEFAULT) # ~45 MB on first run
net.fc = nn.Linear(net.fc.in_features, 10) # new 10-class head
net.fcLinear(in_features=512, out_features=10, bias=True)backbone = tf.keras.applications.MobileNetV2( # ~9 MB on first run
weights='imagenet', include_top=False, input_shape=(160, 160, 3),
pooling='avg')
net = tf.keras.Sequential([backbone, tf.keras.layers.Dense(10, name='head')])
net(tf.zeros((1, 160, 160, 3))) # build the new 10-class head
net.layers[-1]<Dense name=head, built=True>pretrained_net = ResNet50.from_pretrained()
pretrained_net.fc = nnx.Linear(2048, 10, rngs=nnx.Rngs(2))
print('new pretrained-model head:', pretrained_net.fc.kernel.shape)
class Classifier(nnx.Module):
def __init__(self, rngs=None):
rngs = nnx.Rngs(0) if rngs is None else rngs
self.hidden = nnx.Linear(20, 256, rngs=rngs)
self.head = nnx.Linear(256, 2, rngs=rngs) # new head, new name
def __call__(self, x):
return self.head(nnx.relu(self.hidden(x)))
classifier = Classifier()
new_params = nnx.state(classifier, nnx.Param)
[(path, tuple(value.shape)) for path, value in new_params.flat_state()]Downloading bytes: | 0.00B Reconstructing (incomplete total...): | | 0.00B / 0.00B Fetching 1 files: 0%| | 0/1 [00:00<?, ?it/s]new pretrained-model head: (2048, 10)[(('head', 'bias'), (2,)),
(('head', 'kernel'), (256, 2)),
(('hidden', 'bias'), (256,)),
(('hidden', 'kernel'), (20, 256))]net = gluon.model_zoo.vision.resnet18_v2(pretrained=True) # ~42 MB on first run
net.output = nn.Dense(10, in_units=512) # new 10-class head
net.output.initialize()
net.outputDense(512 -> 10, linear)A state_dict is an ordinary Python dict, so adapting one is ordinary dict surgery. We drop the pretrained 1000-class head (we just replaced it) and, to show what a damaged file looks like, also drop one residual block. Loading with strict=False then returns a report of what did not line up instead of raising.
Keras loads weight files whole, so the partial-loading control is a flag rather than dict surgery: load_weights(path, skip_mismatch=True) fills every variable whose saved shape matches and skips the rest, reporting the skips as a warning. To stage a mismatch, save the weights of a donor whose head has 101 classes, our stand-in for a fine-tuned model someone else published, then load that file into the 10-class network. The cell records the warning so we can read it back as a report.
There is no strict=False to lean on: the merge is yours to write, and so is the report. Take from the file every entry whose name and shape match the new model, keep the fresh initialization for the rest, and compute the two key sets that say what happened: missing, parameters the model has but the file did not fill, and unexpected, file entries with no home in the model.
The partial-loading flags exist, but they report nothing: allow_missing=True skips parameters the file did not fill and ignore_extra=True skips file entries with no home in the model, both silently. So the report is yours to compute, the same hand-rolled key-set diff the jax tab writes. Stage a damaged file first: take the pretrained parameters, drop the head we just replaced and, to show what damage looks like, the deepest residual stage (features.8), and save what remains as a plain parameter dict.
pretrained = ResNet18_Weights.DEFAULT.get_state_dict(progress=False)
pretrained = {k: v for k, v in pretrained.items()
if not k.startswith('fc.') and not k.startswith('layer4.')}
damaged_target = resnet18(weights=None)
damaged_target.fc = nn.Linear(damaged_target.fc.in_features, 10)
report = damaged_target.load_state_dict(pretrained, strict=False)
print('missing by block:', dict(Counter(k.split('.')[0]
for k in report.missing_keys)))
print('unexpected:', report.unexpected_keys)missing by block: {'layer4': 25, 'fc': 2}
unexpected: []donor = tf.keras.Sequential([backbone, tf.keras.layers.Dense(101, name='head')])
donor(tf.zeros((1, 160, 160, 3)))
donor.save_weights('donor-101.weights.h5') # stand-in for a downloaded file
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter('always')
net.load_weights('donor-101.weights.h5', skip_mismatch=True)
report = str(caught[-1].message).splitlines()
print(report[0])
print(report[2].split(' Target variable:')[0])A total of 1 objects could not be loaded. Example error message for object <Dense name=head, built=True>:
The shape of the target variable and the shape of the target value in `variable.assign(value)` must match. variable.shape=(1280, 10), Received: value.shape=(1280, 101).file_flat = load_file('mlp-jax.safetensors')
new_flat = flatten(new_params)
matched = {k: v for k, v in file_flat.items()
if k in new_flat and v.shape == new_flat[k].shape}
merged = unflatten({**new_flat, **matched})
print('missing:', sorted(set(new_flat) - set(matched)))
print('unexpected:', sorted(set(file_flat) - set(new_flat)))
nnx.replace_by_pure_dict(new_params, merged)
nnx.update(classifier, new_params)
print('loaded trunk:', jnp.array_equal(
classifier.hidden.kernel, restored['hidden']['kernel']))missing: ['head.bias', 'head.kernel']
unexpected: ['output.bias', 'output.kernel']
loaded trunk: Truedonor = gluon.model_zoo.vision.resnet18_v2(pretrained=True)
file_params = {name: p.data() for name, p in donor.collect_params().items()
if not name.startswith(('output', 'features.8'))}
npx.savez('resnet18-partial.params', **file_params)
file_keys = set(npx.load('resnet18-partial.params'))
damaged_target = gluon.model_zoo.vision.resnet18_v2(classes=10)
damaged_target.initialize()
model_keys = set(damaged_target.collect_params())
damaged_target.load_parameters(
'resnet18-partial.params', allow_missing=True,
ignore_extra=True) # both skip silently
print('missing by block:', dict(Counter('.'.join(k.split('.')[:2])
for k in sorted(model_keys - file_keys))))
print('unexpected:', sorted(file_keys - model_keys))missing by block: {'features.8': 21, 'output.bias': 1, 'output.weight': 1}
unexpected: []Read this report; do not discard it. We loaded the damaged file into a fresh target so every unexplained missing key really remains random. missing_keys lists parameters the model has but the file did not fill. The two fc entries are expected: that head is new and meant to start random. The layer4 entries are a red flag, a whole block of the backbone left uninitialized, which here means the incoming file was incomplete and would produce nonsense features. unexpected_keys, empty here, would list names in the file with no home in the model, the usual sign of a renamed layer. The rule is to name which keys you expect to be missing and treat anything else as a bug.
Read the warning; do not silence it. It names exactly one object that could not be loaded, the head layer, and gives the two shapes that disagree, (1280, 10) against (1280, 101). That is the expected mismatch: the head is new and meant to start random. Any other layer in that list would mean the backbones disagree, a wrong input shape or a renamed layer, and is a bug rather than something to skip. One piece of API drift to know: older tutorials pass by_name=True for partial loads, but in Keras 3 that flag only applies to legacy .h5 files and raises on the native .weights.h5 format; skip_mismatch is the current control.
Read both sets before trusting the merged tree. The two head entries under missing are expected: that layer is new and meant to start random. The two output entries under unexpected are the file’s old head, stranded by the rename, and a renamed layer is what unexpected usually means in practice. The lesson is the same as with a built-in report, only stricter, because nothing prints it unless you do: name which keys you expect in each set and treat anything else as a bug.
Read the diff; nothing else will print it for you. We used a fresh target so unexplained missing keys retain random initialization. The two output entries under missing are expected: that head is new and meant to start random. The 21 features.8 entries are a red flag, a whole residual stage the file failed to deliver; had this file been your only source, those layers would keep whatever values they started with and the features coming out of them would be nonsense. Unexpected, empty here, would list file entries with no home in the model, the usual sign of a renamed layer. The rule is strict because the flags stay silent: name which keys you expect in each set and treat anything else as a bug.
for p in net.parameters():
p.requires_grad = False
for p in net.fc.parameters():
p.requires_grad = True
trainable = sum(p.numel() for p in net.parameters() if p.requires_grad)
total = sum(p.numel() for p in net.parameters())
f'{trainable} trainable of {total}''5130 trainable of 11181642'backbone.trainable = False
trainable = sum(int(tf.size(v)) for v in net.trainable_variables)
total = sum(int(tf.size(v)) for v in net.weights)
f'{trainable} trainable of {total}''12810 trainable of 2270794'def labels(params):
return jax.tree_util.tree_map_with_path(
lambda path, _: ('train' if any(
getattr(p, 'key', None) == 'head' for p in path) else 'freeze'),
params)
tx = optax.multi_transform(
{'train': optax.adam(0.05), 'freeze': optax.set_to_zero()}, labels)
optimizer = nnx.Optimizer(classifier, tx, wrt=nnx.Param)
hidden_before = classifier.hidden.kernel[...].copy()
head_before = classifier.head.kernel[...].copy()
@nnx.jit
def frozen_step(model, optimizer, X):
loss, grads = nnx.value_and_grad(
lambda m: jnp.square(m(X)).mean())(model)
optimizer.update(model, grads)
return loss
_ = frozen_step(classifier, optimizer, X)
print('trunk unchanged:', jnp.array_equal(
hidden_before, classifier.hidden.kernel))
print('head changed:', not jnp.array_equal(head_before, classifier.head.kernel))
sizes = flatten(jax.tree_util.tree_map(jnp.size, merged))
flat_labels = flatten(labels(merged))
trainable = sum(int(s) for k, s in sizes.items() if flat_labels[k] == 'train')
total = sum(int(s) for s in sizes.values())
f'{trainable} trainable of {total}'trunk unchanged: True
head changed: True'514 trainable of 5890'for p in net.collect_params().values():
p.grad_req = 'null'
for p in net.output.collect_params().values():
p.grad_req = 'write'
trainable = sum(p.data().size for p in net.collect_params().values()
if p.grad_req == 'write')
total = sum(p.data().size for p in net.collect_params().values())
f'{trainable} trainable of {total}''5130 trainable of 11187926'torchvision is one source; the Hugging Face Hub is the ecosystem-scale one, and it distributes its weights as safetensors, which closes the loop with the format of the previous section. This section covers how to load and adapt pretrained weights; Section 19.2 covers when it helps and how far to unfreeze.
keras.applications is one source; the Hugging Face Hub is the ecosystem-scale one, and it distributes its weights as safetensors, which closes the loop with the format of the previous section. This section covers how to load and adapt pretrained weights; Section 19.2 covers when it helps and how far to unfreeze.
The Hugging Face Hub distributes JAX weights as safetensors, so the flat dict you just merged has the same shape as the artifact you will download in practice. This section covers how to load and adapt pretrained weights; Section 19.2 covers when it helps and how far to unfreeze.
gluon.model_zoo is one source, frozen where the project stopped; the Hugging Face Hub is the ecosystem-scale one, and it distributes weights as safetensors, which for MXNet means the numpy bridge of the previous section is also your import path. This section covers how to load and adapt pretrained weights; Section 19.2 covers when it helps and how far to unfreeze.
5.6.5 Summary
A saved model is state, not code: a state_dict of tensors that means something only once the code that built the network runs again. For your own files torch.save is fine; for files you share, safetensors stores the same tensors with no executable pickle, which is why hubs standardize on it. A resumable checkpoint bundles more than weights: optimizer state, RNG state, step, and config, written atomically, or a resume restarts the optimizer’s momentum from zero. Loading someone else’s weights is dict surgery plus strict=False, and the missing/unexpected report is a diagnostic to read rather than a warning to silence.
A saved model is state, not code: a collection of path-named variables that means something only once the code that built the network runs again. For your own files np.save and .weights.h5 are fine; for files you share, safetensors stores the same tensors behind a plain JSON header, with no pickle to execute, which is why hubs standardize on it. A resumable checkpoint is a tf.train.Checkpoint of model, optimizer, and step saved as one unit, or a resume restarts the optimizer’s momentum from zero; a fresh optimizer must be built before the restore so Adam’s moments have somewhere to land. Loading someone else’s weights is load_weights with skip_mismatch=True, and the skip warning is a diagnostic to read rather than a message to silence.
A saved model is state, not code: a pytree of named arrays that means something only once the code that built the network runs again. For your own files jnp.save is fine; for files you share, safetensors stores the same tensors behind a plain JSON header, with no pickle to execute, which is why hubs standardize on it. A resumable checkpoint contains model state, optimizer state, and step in pytrees that Orbax saves atomically in a single call, or a resume restarts the optimizer’s momentum from zero. Loading someone else’s weights is pytree surgery, and the missing/unexpected sets you compute are a diagnostic to read rather than a formality to skip.
A saved model is state, not code: a dictionary of dotted-path names to arrays that means something only once the code that built the network runs again. For your own files save_parameters is fine, and its format is pure array data with no pickle to execute; for files you share, safetensors is broadly readable, and the safetensors.numpy bridge is all MXNet needs to speak it. A resumable checkpoint is three versioned files whose completion should be published by one atomic manifest update: parameters, trainer states (Adam’s moments, stored with pickle, so keep such files your own), and a JSON sidecar with the step and config; the RNG stream cannot be snapshotted, only reseeded. Loading someone else’s weights is dict surgery plus allow_missing/ignore_extra, and because those flags skip silently, the missing/unexpected key-set diff is yours to compute and read.
5.6.6 Exercises
- Even if you never deploy to another machine, name two reasons to checkpoint. Then consider the atomic write: if the checkpoint were written straight to its final path (delete the
os.replacefromsave_checkpoint, the rename-into-place that orbax performs, or the fresh numbered files thattf.train.Checkpoint.savewrites), describe the failure a crash mid-write now causes, and why the atomic version avoids it. - Read the first 8 bytes of the safetensors file you wrote for the MLP as a little-endian integer, as the header cell does. How large is the JSON header for the MLP, and how does it grow if you double the hidden width?
- Save the MLP’s parameters cast to
bfloat16and load them back into afloat32model (Section 5.5). What is lost? Is that acceptable for inference? For resuming training? - Take two checkpoints of the regressor 50 steps apart, average their weight tensors into a third set of parameters, and evaluate it. The result previews weight averaging.