How to Load PyTorch Model Parameters into MindSpore

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

Problem Overview

In day-to-day model development and training, many existing open-source projects and paper implementations use PyTorch for model design, development, training, and inference. When we need to develop models with MindSpore, two problems arise:

The model is implemented in PyTorch.
Parameters saved after training a PyTorch model cannot be loaded directly by a MindSpore model.

The first problem can be addressed using the official MindSpore documentation: Typical Differences from PyTorch and PyTorch–MindSpore API Mapping to migrate the model.

For parameter conversion, MindConverter is no longer supported in the latest MindSpore version discussed here. We can therefore convert parameters manually, transforming PyTorch model parameters into a format MindSpore can recognize before loading them.

Solution

I will not repeat the model code migration process here.

The main steps for parameter conversion are:

Load the PyTorch model with PyTorch and obtain its parameters, prams_torch.
Load the MindSpore model with MindSpore and obtain its parameters, prams_ms.
Match PyTorch parameter names to MindSpore parameter names one by one where corresponding parameters exist.
Build a torch_2_ms key mapping and use it to place PyTorch parameter values under the corresponding MindSpore parameter names.
Load the parameters with MindSpore.

Case Study

Different models contain different modules and parameter types. Here, one network illustrates the basic conversion approach; the same reasoning applies to other models.

EfficientNet is a paper published by Google in 2019. See the paper for the detailed network architecture. Here we use EfficientNet+FC as an example of a model with a fully connected layer to explore parameter conversion.

Load the PyTorch Model and Obtain 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')

After this step, the PyTorch model parameters have been saved to param_torch.csv. Inspect the data:

	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

Load the MindSpore Model and Obtain 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')

After this step, the MindSpore model parameters have been saved to prams_ms.csv. Inspect the data:

	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

Match PyTorch Parameter Names to MindSpore Parameter Names

We now have parameter key tables for MindSpore and PyTorch, provided in the attachments. Comparing their naming conventions reveals consistent patterns, including:

Batch Normalization:
- Weights: weight|bias → gamma|beta.
- Moving mean and variance: running_mean|running_var → moving_mean|moving_variance.
Custom blocks: PyTorch uses the _blocks. prefix.
Other differences.

Key Mapping Table

We can use these patterns to write a Python script that converts key names and generates a mapping table:

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

Next, retrieve each weight value from the PyTorch weight dictionary using the corresponding Pytorch_key in the mapping file, wrap it with mindspore.Parameter, and assign it to the corresponding 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

Load the Parameters with MindSpore

load_param_into_net(mindspore_model, ms_values_dict)

The parameters should now be accepted by MindSpore.

What’s more

When storing parameter values, pay attention to differences in parameter precision between PyTorch and MindSpore.

(End)

Model Migration · Tommy Cheese