IX. Model Saving, Loading & Deployment (模型保存、加载与部署)#

1. `torch.save()` / `torch.load()`#

Serializes / deserializes any Python object (model, tensor, dict) to/from a file.

1
# Recommended: save only the weights dict
2
torch.save(model.state_dict(), 'model_weights.pth')
3

4
# Load
5
state = torch.load('model_weights.pth', map_location='cpu')
6
model.load_state_dict(state)

Note: Saving the entire model object couples code paths. Always save only state_dict.

2. `model.state_dict()` / `load_state_dict()`#

Gets / loads an ordered dictionary of model parameters. The core interface for Transfer Learning (迁移学习) and checkpoint resuming (断点续训).

1
checkpoint = {
2
    'epoch': epoch,
3
    'model': model.state_dict(),
4
    'optim': optimizer.state_dict(),
5
    'loss': best_loss
6
}
7
torch.save(checkpoint, 'ckpt.pth')

Note: strict=False allows partial loading (skips missing keys) — common in transfer learning.

3. `torch.jit.script()` / `torch.jit.trace()`#

Compiles a model to TorchScript for deployment in Python-free environments (C++, mobile).

1
# trace: follow execution path (no control flow)
2
traced = torch.jit.trace(model, torch.rand(1, 3, 224, 224))
3
traced.save('traced.pt')
4

5
# script: supports dynamic control flow
6
scripted = torch.jit.script(model)

Note: Models with if/for branches → use script. Pure forward-pass models → use trace (faster).

4. `torch.onnx.export()`#

Exports a PyTorch model to ONNX format for cross-framework deployment (TensorRT, OpenVINO).

1
torch.onnx.export(
2
    model, torch.rand(1, 3, 224, 224), 'model.onnx',
3
    opset_version=17,
4
    input_names=['input'], output_names=['output'],
5
    dynamic_axes={'input': {0: 'batch'}}
6
)

Note: dynamic_axes enables Dynamic Batch Size (动态批次大小) — essential for production deployment. Validate with onnxruntime.

5. `model.parameters()` / `named_parameters()`#

Iterates over all learnable parameters. The named_ version also returns parameter names.

1
total = sum(p.numel() for p in model.parameters() if p.requires_grad)
2
print(f'Params: {total/1e6:.1f}M')
3

4
for name, p in model.named_parameters():
5
    print(name, p.shape)

Note: Set layer-specific learning rates by passing [{'params': p, 'lr': lr}] list to the optimizer.

6. `model.train()` / `model.eval()`#

Switches between training and evaluation modes — affects Dropout and BatchNorm behavior.

1
model.train()
2
for x, y in train_loader:
3
    loss = criterion(model(x), y)
4
    loss.backward(); optimizer.step()
5

6
model.eval()
7
with torch.no_grad():
8
    for x, y in val_loader:
9
        pred = model(x)

Note: Forgetting model.eval() is the most common reason for unstable inference results.

💡 One-line Takeaway
Always save state_dict (not the model object), and remember the deploy path: PyTorch → TorchScript / ONNX → Runtime.

IX. Model Saving, Loading & Deployment (模型保存、加载与部署)#

1. torch.save() / torch.load()#

2. model.state_dict() / load_state_dict()#

3. torch.jit.script() / torch.jit.trace()#

4. torch.onnx.export()#

5. model.parameters() / named_parameters()#

6. model.train() / model.eval()#