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

問題の概要

日々のモデル開発や学習では、既存のオープンソースプロジェクトや論文の再実装の多くが、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に保存されます。データを確認しましょう：

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に保存されます。データを確認しましょう：

PytorchとMindSporeのパラメータ名を一対一に対応付ける

これでMindSporeとPytorchそれぞれのパラメータのキー名一覧が得られました。添付ファイル欄にも掲載しています。両者の命名の違いを観察すると、一定の規則が見つかります。たとえば次のとおりです：

• Batch Normalization：
  • 重み：weight|bias——gamma|beta。
  • 移動平均と分散：running_mean|running_var——moving_mean|moving_variance。
• 独自のblocks：pytorchには_blocks.という接頭辞が付いています。
• その他

キー名の対応表

この規則に従ってPythonスクリプトを書けば、キー名を変換し、対応表を生成できます：

続いて、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の数値精度の違いに注意してください。

（完）