1、模型保存和加载
主要有两种情况:一是仅保存参数,二是保存参数及模型结构。
保存参数:
torch.save(net.state_dict())
加载参数(加载参数前需要先实例化模型):
param = torch.load('param.pth')
net.load_state_dict(param)
保存模型结构和参数:
torch.save(net)
加载模型:
net = torch.load('model.pt')
2、解析模型权重文件
当加载某个模型文件后,如果需要查看模型中的算子和参数,可以将模型解析为字典,然后逐一打印。
以lent5为例,将lenet5模型保存为权重文件,然后重新加载权重文件并解析其中每一层的参数。
参考代码:
def pytorch_params(pth_file):
par_dict = torch.load(pth_file, map_location='cpu')
for name in par_dict:
parameter = par_dict[name]
print(name, parameter.numpy().shape)
以上代码是加载的权重文件,文件只有参数,没有模型结构,如果加载的是包含模型结构的权重文件,可以做如下修改:
def pytorch_params(pt_file):
net = torch.load(pt_file, map_location='cpu')
par_dict = net.state_dict()
for name in par_dict:
parameter = par_dict[name]
print(name, parameter.numpy().shape)
解析结果:
3、加载自定义参数
某些情况下可能需要对某个算子进行单独调试,如加载特定参数进行推理计算,用来确定输出结果符合预期。以Conv2d算子为例进行测试,首先设定卷积层输入为3,输出为3,卷积核为3*3,偏置bias为False。通过numpy随机一个3*3*3*3的矩阵作为自定义参数,将参数转换为Tensor以后,添加到dict中,然后通过load_state_dict将参数加载进网络。
参考脚本:
import torch
import torch.nn as nn
import numpy as np
net = nn.Conv2d(3, 3, kernel_size=(3, 3), padding=1, bias=False)
param = np.random.random((3, 3, 3, 3))
param = param.astype(np.float32)
torch_param = {'weight': torch.Tensor(param)}
net.load_state_dict(torch_param)
net.eval()
data = np.random.random((1, 3, 16, 16))
data = data.astype(np.float32)
result = net(torch.Tensor(data))
print(result)