實用分享：如何將 PyTorch 模型參數載入 MindSpore 模型？

Fri, 01 Sep 2023 22:53:58 +0530

問題簡述

在日常的模型開發、訓練過程中我們經常會遇到這樣的現象：在現有的開源專案或者論文復現中，多數模型使用Pytorch設計、開發和訓練推理，當我們需要使用MindSpore框架進行模型開發時，會遇到以下兩個問題：

模型使用Pytorch編碼；
Pytorch模型訓練後儲存的引數無法被MindSpore模型直接載入。

對於第一個問題，我們可以根據昇思官方提供的文件：與Pytorch典型區別和PyTorch與MindSpore API對映表來完成模型的遷移；

對於模型引數的轉換，在最新的MindSpore版本中MindConverter不再支援，因此可以考慮針對模型引數，我們進行手動的轉換，將Pytorch模型引數轉換為MindSpore能識別的格式後，再進行載入。

解決方案

模型的編碼轉換不再贅述。

引數轉換主要思路如下：

使用Pytorch載入Pytorch模型，並取得模型引數prams_torch；
使用MindSpore載入MindSpore模型，並取得模型引數prams_ms；
將Pytorch模型的引數名和MindSpore模型引數名一一對應（有的話）；
建立torch_2_ms鍵名對映表，使用鍵名對映表將Pytorch模型引數值載入到MindSpore引數名對應的位置上；
使用MindSpore載入引數。

案例分析

不同模型的模組不相同，引數型別也不盡相同，此處我們以一個網路舉例，說明轉換的基本思路，不同的模型其轉換思路是類似的。

EfficientNet是谷歌於2019年發表的文章，詳細網路架構可檢視文章描述，此處我們以EfficientNet+FC全連線層的模型為例，探討如何進行網路模型引數的轉換。

使用Pytorch載入Pytorch模型，並取得模型引數prams_torch

import torch
from test.efficientnet_pytorch.model import EfficientNet as EN_pytorch
import pandas as pd

pytorch_model = EN_pytorch.from_name(cfg['model'], override_params={'num_classes': 3})
pytorch_model.cuda()
pytorch_weights_dict = pytorch_model.state_dict()
param_torch = pytorch_weights_dict.keys()
param_torch_lst = pd.DataFrame(param_torch)
param_torch_lst.to_csv('param_torch.csv')

步驟結束後，我們就將pytorch的模型引數存到了param_torch.csv下，觀察資料：

	keys
0	_conv_stem.weight
1	_bn0.weight
2	_bn0.bias
3	_bn0.running_mean
4	_bn0.running_var
5	_bn0.num_batches_tracked
6	_blocks.0._depthwise_conv.weight
7	_blocks.0._bn1.weight
8	_blocks.0._bn1.bias
9	_blocks.0._bn1.running_mean
10	_blocks.0._bn1.running_var

使用MindSpore載入MindSpore模型，並取得模型引數prams_ms

import mindspore as ms
from test.efficientnet_mindspore.model import EfficientNet as EN_ms
import pandas as pd

mindspore_model = EN_ms.from_name(cfg['model'], override_params={'num_classes': 3})
prams_ms = mindspore_model.parameters_dict().keys()
prams_ms_lst = pd.DataFrame(prams_ms)
prams_ms_lst.to_csv('prams_ms.csv')

步驟結束後，我們就將MindSpore的模型引數存到了prams_ms.csv下，觀察資料：

	keys
0	_conv_stem.weight
1	_bn0.moving_mean
2	_bn0.moving_variance
3	_bn0.gamma
4	_bn0.beta
5	0._depthwise_conv.weight
6	0._bn1.moving_mean
7	0._bn1.moving_variance
8	0._bn1.gamma
9	0._bn1.beta
10	0._se_reduce.weight

將Pytorch模型的引數名和MindSpore模型引數名一一對應

自此我們就得到了MindSpore和Pytorch各自的引數鍵名錶（附在附件區域），隨後觀察二者引數命名上的差異，可以發現固定的規律，以下述幾個方面為例：

Batch Normalization：
- 權重：weight|bias——gamma|beta；
- 移動加權和方差：running_mean|running_var——moving_mean|moving_variance；
自定義blocks：pytorch帶前置的_blocks.；
其他

鍵名對映表

這樣就可以根據規律寫出一個Python指令碼來完成鍵名的轉化，並生成鍵名對映表：

Pytorch	mindspore
_conv_stem.weight	_conv_stem.weight
_bn0.weight	_bn0.gamma
_bn0.bias	_bn0.beta
_bn0.running_mean	_bn0.moving_mean
_bn0.running_var	_bn0.moving_variance
_blocks.0._depthwise_conv.weight	0._depthwise_conv.weight
_blocks.0._bn1.weight	0._bn1.gamma
_blocks.0._bn1.bias	0._bn1.beta
_blocks.0._bn1.running_mean	0._bn1.moving_mean
_blocks.0._bn1.running_var	0._bn1.moving_variance
_blocks.0._se_reduce.weight	0._se_reduce.weight

隨後在Pytorch的權重字典中，按照對應檔案的Pytorch_key取出權重值，隨後使用mindspore.Parameter進行封裝，新增到mindspore.key對應的權值中去：

for i in ms_param_lst.values:
    ms_key = i
    pt_key = param_mapping[ms_key]
    pt_val = pt_values_dict[pt_key]
    if not isinstance(pt_val, np.ndarray):
        pt_val = pt_val.cpu().numpy()
    ms_val = Parameter(pt_val, ms_key)
    print(ms_val)
    ms_values_dict[ms_key] = ms_val

使用MindSpore載入引數

load_param_into_net(mindspore_model, ms_values_dict)

此時，引數應該就可以被MindSpore接受了。

What’s more

在引數值的儲存過程中，要注意Pytorch和MindSpore引數精度的差異；

（完）

模型遷移 · Tommy Cheese