PyTorch のモデルパラメーターを MindSpore に読み込む方法

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

問題の概要

日々のモデル開発や学習では、既存のオープンソースプロジェクトや論文の再実装の多くが、Pytorchでモデルを設計、開発し、学習や推論を行っているという状況によく出会います。MindSporeフレームワークを使ってモデルを開発したい場合、次の二つの問題があります：

モデルのコードがPytorchで書かれていること。
Pytorchの学習後に保存されたパラメータを、MindSporeのモデルで直接読み込めないこと。

一つ目の問題は、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は、Googleが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がパラメータを読み込めるはずです。

補足

パラメータ値を保存するときは、PytorchとMindSporeの数値精度の違いに注意してください。

（完）

モデル移行 · Tommy Cheese