Pytorch 网络模型的保存与读取

Posted 2021-08-30 ʚVVcatɞ

保存方式一：

例如：对 vgg16 网络模型进行保存，模型如下

模型保存使用 torch.save() 方法

torch.save(模型, "文件名.h5")  # 保存 模型结构 + 模型参数

使用方式如下：

import torch
import torchvision

vgg16 = torchvision.models.vgg16(pretrained=False)  # 加载vgg16网络模型

# 保存方式1
torch.save(vgg16, "vgg16_method.pth")  # 保存为 vgg16_method.pth 文件

加载保存后的模型文件

import torch

model = torch.load("vgg16_method.pth")
print(model)

保存方式二：

import torch
import torchvision

vgg16 = torchvision.models.vgg16(pretrained=False)  # 加载vgg16网络模型

# 保存模型的参数
torch.save(vgg16.state_dict(), "vgg16_method.pth")  # 把模型的参数 保存成字典

加载保存后的模型文件

import torch

model = torch.load("vgg16_method.pth") # 使用 torch.load 加载 模型文件，得到模型的参数
print(model)

使用以下方法得到完整的网络模型

import torch
import torchvision

vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_method.pth"))
print(vgg16)

保存自定义模型：

下面是自定义了一个 VVcatModel 模型进行保存。

import torch
from torch import nn

class VVcatModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

    def forward(self, x):
        x = self.conv1(x)
        return x

vvcat = VVcatModel() # 实例化模型
torch.save(vvcat, "vgg16_method.pth")

加载自定义模型

import torch
from torch import nn

class VVcatModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

    def forward(self, x):
        x = self.conv1(x)
        return x
        
# 此处不需要实例化模型
model = torch.load("vgg16_method.pth")
print(model)

注：自定义保存模型与加载自定义模型区别在于，加载模型时，不需要实例化模型。

注意：
如果使用GPU去训练生成的网络模型文件，在加载时可能出现以下错误

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

当使用GPU去训练生成的网络模型文件，在加载时使用的是CPU去加载，所以，在加载该网络文件时，需要指定CPU去加载该文件， map_location=torch.device(‘cpu’)，添加如下：

model = torch.load("./vgg16_method.pth", map_location=torch.device('cpu'))