当保存和加载模型时，需要熟悉三个核心功能：

保存

torch.save(model.state_dict(), PATH)

加载

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

torch.save(model, PATH)

new_model = torch.load(PATH)
new_model.eval()
#new_model 不再需要第一种方法中的建立新模型的步骤

torch.save({
 'epoch': epoch,
 'model_state_dict': model.state_dict(),
 'optimizer_state_dict': optimizer.state_dict(),
 'loss': loss,
  ...
  }, PATH)

model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)
checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
model.eval()
# - or -
model.train()

torch.save({
 'modelA_state_dict': modelA.state_dict(),
 'modelB_state_dict': modelB.state_dict(),
 'optimizerA_state_dict': optimizerA.state_dict(),
 'optimizerB_state_dict': optimizerB.state_dict(),
 ...
 }, PATH)

modelA = TheModelAClass(*args, **kwargs)
modelB = TheModelBClass(*args, **kwargs)
optimizerA = TheOptimizerAClass(*args, **kwargs)
optimizerB = TheOptimizerBClass(*args, **kwargs)
checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])
modelA.eval()
modelB.eval()
# - or -
modelA.train()
modelB.train()

torch.save(modelA.state_dict(), PATH)

modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)

torch.save(model.state_dict(), PATH)

device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))

torch.save(model.state_dict(), PATH)

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # 确保在你提供给模型的任何输入张量上调用input = input.to(device)

torch.save(model.state_dict(), PATH)

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0")) # Choose whatever GPU device number you want
model.to(device) # 确保在你提供给模型的任何输入张量上调用input = input.to(device)

torch.save(model.module.state_dict(), PATH)

# 加载任何你想要的设备

Pytorch目前成为学术界最流行的DL框架，没有之一。很大程度上，简洁直观地操作有关。模型的保存和加载，于pytorch而言，也是很简单的。本文做了一个比较实验，方便大家理解。
首先，要清楚几个函数：torch.save，torch.load，state_dict()，load_state_dict()。
先举最简单的例子：

import torch

model = torch.load('my_model.pth')
torch.save(model, 'new_model.pth')

上面的代码非常直观，一载一存。但是有一个问题，这样保存的pth文件直接包含了整个模型的结构。当你需要灵活加载模型参数时，比如只加载部分参数，那么这种情况保存的pth文件读取进来还得额外解析出“参数文件”。

如果想更灵活对待咱们训练好的模型参数，咱们可以使用下面这个方法。pytorch把所有的模型参数用一个内部定义的dict进行保存，自称为“state_dict”。这个所谓的state_dict就是不带模型结构的模型参数了~
咱们的加载和保存就发生了一点微妙的变化：

import torch
model = MyModel() # init your model class, build the graph shape

state_dict = torch.load('model_state_dict.pth')
model.load_state_dict(state_dict)

torch.save(model.state_dict(), 'model_state_dict1.pth')

比较上面两段代码，咱们可以有一下结论：

pth文件既可能保存了模型的图结构，也有可能没保存；
加载没保存图结构的pth时，需要先初始化模型结构，即把架子搭好；
在保存模型的时候，如果不想保存图结构，可以单独保存model.state_dict()

• 实验

import torch
import torchvision.models as models

model = models.vgg16(pretrained=True)
torch.save(model.state_dict(), 'only_weights.pth')

model_state_dict = torch.load('only_weights.pth')
model1 = models.vgg16() # describe the graph shape
model1.load_state_dict(model_state_dict)
model1.eval()

torch.save(model1, 'whole_model.pth')

model2 = torch.load('whole_model.pth')
model2.eval()

# model3 = torch.load('only_weights.pth')
# model3.eval()    # Error

model3切换到eval()模式就会报错，原因是model3只包含weights而缺乏图结构~

• torch.load_state_dict()函数的用法

在Pytorch中构建好一个模型后，一般需要进行预训练权重中加载。torch.load_state_dict()函数就是用于将预训练的参数权重加载到新的模型之中，操作方式如下所示：

sd_net = torchvision.models.resnte50(pretrained=False)
sd_net.load_state_dict(torch.load('*.pth'), strict=True)

在本博文中重点关注的是 属性 strict; 当strict=True,要求预训练权重层数的键值与新构建的模型中的权重层数名称完全吻合；如果新构建的模型在层数上进行了部分微调，则上述代码就会报错：说key对应不上。

此时，如果我们采用strict=False 就能够完美的解决这个问题。也即，与训练权重中与新构建网络中匹配层的键值就进行使用，没有的就默认初始化。

参考博文：
https://blog.csdn.net/ChaoMartin/article/details/118686268

https://blog.csdn.net/leviopku/article/details/123925804

pytorch中保存和加载模型是绑在一起的。
这里我需要注意一下不同的保存方式对应不同的读取方式，两者各有利弊。

首先说说pytorch.save()这个函数，可以参考官网：pytroch.save。
简而言之，这个函数可以保存任意的东西，比如tensor或者模型，或者仅仅是模型的参数。
如果将保存对象局限在模型上，通常来说我们有两种方式：直接保存所有的模型，只保存模型中的参数（模型结构就保存了）。以下分别说说两种不同的方式。

为了说明，我们先建立一个简单的模型。

import torch
import torch.nn as nn

class Generator(nn.Module):
    def __init__(self, in_c, out_c, ngf=64):
        super(Generator, self).__init__()
        model = []
        model += [
            nn.Conv2d(in_c, ngf, 3, 2, 1),
            nn.ReLU(),
            nn.BatchNorm2d(ngf),
            nn.Conv2d(ngf, out_c, 3, 2, 1)
        ]
        self.model = nn.Sequential(*model)
        
    def forward(self, x):
        return self.model(x)

netG = Generator(3, 3)
input = torch.zeros(10, 3, 256, 256)
output = netG(input)

直接保存所有模型并读取

直接使用简单粗暴的方式保存：

torch.save(netG, 'netG.pt')

对应的，我们可以这样读取模型

netC = torch.load('netG.pt')
input = torch.zeros(10, 3, 256, 256)
output = netC(input)

正常情况如下（警告先忽略）：

只保存模型中的参数并读取

我们说模型的参数保存在网络的state_dict中，使用这个就可以读取网络的参数了。

torch.save({'netG': netG.state_dict()}, 'model_test.pt')

对应的加载模型的方式如下：

netD = Generator(3, 3)
state_dict = torch.load('model_test.pt')
netD.load_state_dict(state_dict['netG'])
input = torch.zeros(10, 3, 256, 256)
output = netD(input)

总结

我们可以看到第一种方法可以直接保存模型，加载模型的时候直接把读取的模型给一个参数就行。而第二种方法则只是保存参数，在读取模型参数前要先定义一个模型（模型必须与原模型相同的构造），然后对这个模型导入参数。虽然麻烦，但是可以同时保存多个模型的参数，而第一种方法则不能，而且第一种方法有时不能保证模型的相同性（你读取的模型并不是你想要的）。

总的来说，我们一般来选择第二种来保存和读取。
退一步讲，如何保存模型决定了如何读取模型。