第五章模型篇：模型保存与加载

最新推荐文章于 2025-08-13 19:13:08 发布

原创

最新推荐文章于 2025-08-13 19:13:08 发布 · 2.1k 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #pytorch #python

本文介绍了PyTorch中模型的保存与加载机制，包括torch.save()和torch.load()函数的使用，以及state_dict的保存和加载。模型的参数可以通过state_dict保存，然后加载到具有相同结构的新模型中。此外，文章还讨论了保存整个模型、checkpoint的概念，以及如何保存和加载多个模型的状态。

参考教程：
https://pytorch.org/tutorials/beginner/basics/saveloadrun_tutorial.html

训练好的模型，可以保存下来，用于后续的预测或者训练过程的重启。
为了便于理解模型保存和加载的过程，我们定义一个简单的小模型作为例子，进行后续的讲解。

这个模型里面包含一个名为self.p1的Parameter和一个名为conv1的卷积层。我们没有给模型定义forward()函数，是因为暂时不需要用到该方法。假如你想使用这个模型对数据进行前向传播，会返回 “NotImplementedError: Module [Model] is missing the required “forward” function”。

import torch
import torch.nn as nn
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.t1 = torch.randn((3,2))
        self.p1 = nn.Parameter(self.t1)
        self.conv1 = nn.Conv2d(1, 1, 5)
net = Model()

pytorch中的保存与加载

首先我们来看一下pytorch中的保存和加载的方法是怎么实现的。

torch.save()

参考文档：https://pytorch.org/docs/stable/generated/torch.save.html
首先来看一下torch.save()函数。

torch.save(obj, f, pickle_module=pickle, pickle_protocol=DEFAULT_PROTOCOL, _use_new_zipfile_serialization=True)

torch.save()函数传入的第一个参数，就是我们要保存的对象，它的类别要求是object，而没有限定在nn.Module()或者nn.Parameters()等等之间。说明它可以保存的类型是多种多样的，很灵活。
传入的第二个参数是f，f是一个file-like object或者文件路径，也就是我们想要保存的位置。
后面的几个参数可以不用管它，一般也不会用到。从参数名称可以看到，我们想要保存的object是以pickle的形式保存的。因为pickle支持多种数据类型。
在源码中给了两个使用torch.save的例子。

  >>> # xdoctest: +SKIP("makes cwd dirty")
        >>> # Save to file
        >>>

最低0.47元/天解锁文章