Pytorch学习日记7：Save and Load the Model

主要内容：

1.保存载入模型权重

2.保存和载入模型形状

一、保存载入模型权重

保存模型权重。PyTorch模型将学到的参数存储在内部状态字典中，称为state_dict。这些可以通过torch.save方法持久化。直接上代码：

import torch
import torchvision.models as models
model = models.vgg16(weights='IMAGENET1K_V1')
torch.save(model.state_dict(), 'model_weights.pth')

torchvision.models包含用于处理不同任务的模型定义。这段代码是使用PyTorch库加载预训练的VGG16模型（图像分类的模型结构），并将其权重保存到名为'model_weights.pth'的文件中。

1. 导入PyTorch库中的models模块。
2. 使用models.vgg16()函数创建一个预训练的VGG16模型实例。
3. 设置模型参数weights='IMAGENET1K_V1'，表示使用在ImageNet数据集上预训练的权重。
4. 使用torch.save()函数将模型的状态字典（state_dict）保存到文件'model_weights.pth'中。

为了加载模型的权重，你需要先创建一个相同模型的实例，然后用load_state_dict()方法加载参数。直接上代码：

model = models.vgg16() # we do not specify "weights", i.e. create untrained model
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

这段代码是使用PyTorch库加载预训练的VGG16模型。首先，它创建了一个未指定权重的VGG16模型，然后从名为'model_weights.pth'的文件中加载预训练的权重。最后，它将加载的权重应用到模型中。

在预测前一定要先调用model.eval()方法来设置dropout和batch normalization层为评估模型，否则会导致不一致的预测结果。留一个坑，还不太理解这个地方！

二、保存载入模型形状

在加载模型权重时，我们需要先将模型类实例化，因为该类定义了网络的结构。我们可能想把这个类的结构和模型一起保存，在这种情况下，我们可以把model（而不是model.state_dict()）传给保存函数。

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

加载模型如下：

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