Pytorch学习笔记十七：模型的保存与加载

Dear_林

已于 2022-05-07 11:21:19 修改

阅读量2.9k

点赞数 4

分类专栏： pytorch 文章标签： pytorch 学习深度学习

于 2022-05-06 11:46:00 首次发布

本文链接：https://blog.csdn.net/dear_learner/article/details/124562200

版权

本文详细介绍了PyTorch中模型的保存与加载，包括序列化与反序列化、模型的完整保存与参数保存，以及如何进行断点训练。此外，还讨论了模型的微调（finetune）以及GPU的使用，如数据迁移、多GPU并行运算及常见错误处理。

摘要由CSDN通过智能技术生成

一、模型的保存与加载

当我们的模型训练好之后是需要保存下来，以备后续的使用，那么如何保存和加载模型呢？下面就从三个方面来理解一下。

1、序列化与反序列化

序列化是指内存中的某一对象保存到硬盘中，以二进制的形式存储下来，这就是一个序列化的过程；反序列化就是将硬盘中的存储的二进制数反序列化到内存中，得到一个相应的对象，这样就可以再次使用这个模型了。如下图所示：
在这里插入图片描述
序列化和反序列化的目的就是将模型保存并再次使用。

pytorch中序列化和反序列化的方法：
torch.save(obj, f)：obj表示对象，也就是保存的数据，可以是模型、张量、dict等。f表示输出的路径。
torch.load(f, map_location)：f表示文件的路径，map_location 表示指定存放的位置，CPU或者GPU

2、模型的保存与加载

pytorch中模型的保存有两种方式，一种是保存整个Module，另一种是保存模型的参数，相对应的加载模型也就有两种方式。
1⃣️保存和加载整个Module：torch.save(net, path)，torch.load(fpath)
2⃣️保存模型参数：torch.save(net.state_dict(), path), net.load_state_dict(torch.load(path))

第一种方法保存的是整个模型架构，比较费时占用内存，第二种方法只保存了模型上可学习的参数，等建立一个新的相同的网络结构时将这些参数加载到网络中即可。推荐使用第二种。
下面从代码中看一下这两种方法的具体实现过程：
构建一个网络结构：

class LeNet2(nn.Module):
    def __init__(self, classes):
        super(LeNet2, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 6, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(16*5*5, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size()[0], -1)
        x = self.classifier(x)
        return x

    def initialize(self):
        # 这里将网络结构中的参数都填充为20191104
        for p in self.parameters():
            p.data.fill_(20191104)
# 建立网络结构
net = LetNet2(classes=2019)

# 训练
# "训练"
print("训练前: ", net.features[0].weight[0, ...])
net.initialize