7. SAVE AND LOAD MODEL

7. SAVE AND LOAD MODEL

import torch
import torchvision.models as models

7.1 保存并加载模型权重

PyTorch模型将学习到的参数存储在称为state_dict的内部字典中，并通过torch.save方法持久化：

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

要加载模型权重，需要首先创建同一模型的实例，然后使用load_state_dict()方法加载参数。

model = models.vgg16() 
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

7.2 保存并加载带有形状的模型

当加载模型权重时，需要首先实例化模型类，因为该类定义了网络的结构。我们可能想将这个类的结构与模型一起保存，在这种情况下，我们可以将model传递给保存函数：

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

可以以如下方式加载模型：

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