最近在做一些pytorch代码迁移到mindspore的项目,由于训练数据非常大不好重新训练,所以想把作者提供的预训练模型参数迁移一下看看效果,这篇文章相当于做个笔记。
这里以deeplab的Resnet101为例:
需要的东西
-
pytorch版本的权重文件(.pth文件)
-
mindspore的模型代码
mindspore的模型代码需要根据pytorch版代码来自己迁移,有些比较常用的预训练模型代码可以在mindspore的modelzoo中找到。
首先导入需要的包
import mindspore
import torch
import pandas as pd
import csv
import numpy as np
from mindspore import load_checkpoint,load_param_into_net,save_checkpoint
from graphs.models import deeplab_multi # 导入你的mindspore版本网络代码
读取双方的参数并把键值(key)保存为csv文件
# mindspore模型
ms_model = deeplab_multi.DeeplabMulti(num_classes=21,pretrained=False)
keys = pd.DataFrame(ms_model.parameters_dict().keys())
keys.to_csv('mindspore_keys.csv')
# Pytorch参数
torch_model = torch.load('GTA5_source.pth') # 替换成自己的.pth文件路径
model_key = pd.DataFrame(model.keys())
model_key.to_csv('pytorch_res101.csv')
比较一下两边的区别
双方一对一的key可以很容易辨别,可以查看这篇博客
实用干货:如何把Pytorch模型参数加载到MindSpore模型?-CSDN博客
此外,在pytorch中多了一个
num_batches_tracked
参数,它记录了自从模型训练开始以来处理的批次数量。这个参数不是用于直接参与前向或反向传播的计算,而是作为一个状态量,通常用于指导 BatchNorm 层的行为。
在迁移 PyTorch 模型到 MindSpore 时,可以忽略该参数。
示例模型只采用了conv2d层和BN层,如果有采用其他层需要再甄别一下
制作双方key值一一对应的字典
# 把双方的key值读进来
torch_keys = pd.read_csv('pytorch_res101.csv')
mindspore_keys = pd.read_csv('mindspore_keys.csv')
# 把'num_batches_tracked'参数踢掉
torch_list = []
for index,value in torch_keys.iterrows():
name = value[-1]
if not 'num_batches_tracked' in name:
torch_list.append(name)
# 制作字典
ms_list = []
for index,value in mindspore_keys.iterrows():
name = value[-1]
ms_list.append(name)
key_dict = {}
if len(torch_list) == len(ms_list):
for p in range(len(torch_list)):
key_dict[ms_list[p]] = torch_list[p]
然后把Pytorch参数写入mindspore模型
for k,v in ms_model.parameters_dict().items():
torch_k = key_dict[k]
torch_v = torch_model[torch_k]
if not isinstance(torch_v,np.ndarray):
torch_v = torch_v.cpu().numpy()
ms_v = mindspore.Parameter(torch_v,name=k)
ms_model.parameters_dict()[k] = ms_v
最后把模型保存好就完成了
save_checkpoint(ms_model,'deeplab_multi_ms.ckpt')
另外还有一点:
在pytorch中直接查看某层的参数字典,running_mean 和 running_var 两个参数是不作显示的:
# 单独创建一个bn层
import torch.nn as nn
bn1 = nn.BatchNorm2d(64, momentum=0.1)
params = bn1.named_parameters()
for i in params:
print(i[0])
>>>
weight
bias
但可以通过加后缀的方式来查看:
bn1.running_mean
>>>
tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])