本文共 648 字，大约阅读时间需要 2 分钟。

PyTorch 中的模型保存与加载是机器学习训练过程中的常见操作。以下是两种常用的保存与加载方式：

1. 保存整个模型

这种方法适用于需要保留整个模型结构和训练权重的场景。使用 torch.save() 函数将模型及其状态字典一起保存。

# 保存整个模型及其状态字典torch.save(resnet, 'model.ckpt')

2. 仅保存模型参数

为了更高效地保存和加载模型，只需保存模型的参数状态字典即可。这种方法在模型迭代过程中特别有用，可以减少存储空间和加载时间。

# 仅保存模型参数torch.save(resnet.state_dict(), 'params.ckpt')

3. 加载模型

在需要使用预训练模型时，可以使用 torch.load() 函数加载保存的文件。

加载整个模型

如果使用了 model.ckpt 文件，直接加载整个模型即可。

model = torch.load('model.ckpt')

加载模型参数

如果使用了 params.ckpt 文件，需要同时加载模型结构和参数。

resnet.load_state_dict(torch.load('params.ckpt'))

注意事项

• 文件路径：确保文件路径正确，避免文件名冲突。
• 兼容性：确保加载的模型与当前环境的 PyTorch 版本兼容。
• 错误处理：在加载模型时，若遇到 KeyError，请检查文件是否损坏或路径是否正确。

通过上述方法，您可以轻松地在训练过程中保存模型，并在需要时加载继续使用。

转载地址：http://sgefk.baihongyu.com/