PyTorch深度学习计算：模型与张量的文件I/O操作详解-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00545/article/details/148527330

PyTorch深度学习计算：模型与张量的文件I/O操作详解

d2l-pytorch dsgiitr/d2l-pytorch: d2l-pytorch 是Deep Learning (DL) from Scratch with PyTorch系列教程的配套代码库，通过从零开始构建常见的深度学习模型，帮助用户深入理解PyTorch框架以及深度学习算法的工作原理。项目地址: https://gitcode.com/gh_mirrors/d2/d2l-pytorch

引言

在深度学习项目中，模型训练完成后，我们需要将训练好的模型保存下来以便后续使用或部署。同时，在长时间训练过程中，定期保存中间结果（检查点）也是防止意外中断导致训练成果丢失的重要实践。本文将详细介绍PyTorch框架中张量和模型参数的保存与加载方法，帮助读者掌握深度学习模型持久化的关键技术。

张量的保存与加载

PyTorch提供了简单直接的API来保存和加载张量数据，这是最基本的I/O操作。

单个张量的保存与加载

import torch

# 创建一个张量
x = torch.arange(4, dtype=torch.float32)

# 保存张量到文件
torch.save(x, "x-file")

# 从文件加载张量
x2 = torch.load("x-file")
print(x2)  # 输出: tensor([0., 1., 2., 3.])

多个张量的保存与加载

我们可以将多个张量组织成列表或字典进行保存和加载：

# 保存多个张量
y = torch.zeros(4)
torch.save([x, y], 'x-files')

# 加载多个张量
x2, y2 = torch.load('x-files')

# 使用字典保存张量
mydict = {'x': x, 'y': y}
torch.save(mydict, 'mydict')
mydict2 = torch.load('mydict')

技术要点：

保存的文件格式是PyTorch特有的二进制格式
加载时数据类型和形状会保持原样
这种方法适用于任何PyTorch张量，无论其维度如何

模型参数的保存与加载

与保存简单张量相比，保存整个模型的参数需要更复杂的处理，因为PyTorch模型不仅包含参数，还包含计算图结构。

模型定义与初始化

首先我们定义一个简单的多层感知机(MLP)模型：

from torch import nn

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(20, 256)
        self.output = nn.Linear(256, 10)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        H_1 = self.relu(self.hidden(x))
        return self.output(H_1)

# 实例化模型并进行前向传播
net = MLP()
x = torch.randn(size=(2, 20))
y = net(x)

保存模型参数

PyTorch使用state_dict()方法获取模型的所有参数状态：

torch.save(net.state_dict(), 'mlp.params')

关键理解：

state_dict()返回一个有序字典，包含模型的所有可学习参数
只保存参数，不保存模型结构
文件大小通常比保存整个模型小很多

加载模型参数

要恢复模型，我们需要先实例化相同结构的模型，然后加载参数：

clone = MLP()  # 必须与原始模型结构相同
clone.load_state_dict(torch.load("mlp.params"))
clone.eval()  # 设置为评估模式

# 验证加载的模型
y_clone = clone(x)
print(torch.all(y_clone == y))  # 应该输出True

注意事项：

模型结构必须与保存时完全一致
加载后通常需要调用eval()将模型设为评估模式
如果用于训练，需要确保所有必要的层(如Dropout、BatchNorm等)状态正确

技术深入：模型保存的底层原理

PyTorch的模型保存机制实际上是将模型的state_dict序列化存储。state_dict包含：

每一层的权重参数(weight)
每一层的偏置参数(bias)
优化器状态(如果保存)
其他需要持久化的模型状态

这种设计实现了模型结构与参数的分离，带来了几个优势：

可以灵活地修改模型结构而重用部分参数
参数文件更小，便于传输和存储
支持跨平台部署(需考虑兼容性)

实际应用建议

检查点保存：在长时间训练中定期保存模型状态

# 每epoch结束时保存
torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}, f'checkpoint_epoch{epoch}.pt')

模型部署：保存训练好的模型用于推理

# 保存完整模型(包含结构)
torch.save(model, 'full_model.pt')
# 注意：这种方式可能在不同PyTorch版本间不兼容

参数迁移：将部分层参数迁移到新模型

# 假设我们只需要hidden层的参数
pretrained_dict = torch.load('mlp.params')
model_dict = new_model.state_dict()

# 筛选需要的参数
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
model_dict.update(pretrained_dict)
new_model.load_state_dict(model_dict)