Как загрузить параметры модели PyTorch в MindSpore

Tommy Cheese | · Время чтения: 3 мин
Как загрузить параметры модели PyTorch в MindSpore

Описание задачи

При разработке и обучении моделей часто встречается следующая ситуация: в открытых проектах и воспроизведениях научных работ большинство моделей спроектированы и реализованы на Pytorch, который также используется для обучения и инференса. Если же разработку нужно вести на MindSpore, возникают две проблемы:

  • Код модели написан на Pytorch.
  • Сохранённые после обучения параметры модели Pytorch нельзя напрямую загрузить в модель MindSpore.

Первую проблему можно решить переносом модели по официальной документации MindSpore: Типичные отличия от Pytorch и Таблица соответствия API PyTorch и MindSpore.

Для преобразования параметров MindConverter в рассматриваемой актуальной версии MindSpore больше не поддерживается. Поэтому можно выполнить ручное преобразование: перевести параметры модели Pytorch в формат, распознаваемый MindSpore, и затем загрузить их.

Решение

На преобразовании кода модели подробно останавливаться не будем.

Основная схема преобразования параметров:

  • Загрузить модель Pytorch средствами Pytorch и получить параметры prams_torch.
  • Загрузить модель MindSpore средствами MindSpore и получить параметры prams_ms.
  • Установить взаимно однозначное соответствие между именами параметров моделей Pytorch и MindSpore, где оно существует.
  • Создать таблицу соответствия ключей torch_2_ms и с её помощью поместить значения параметров Pytorch в позиции, соответствующие именам параметров MindSpore.
  • Загрузить параметры средствами MindSpore.

Разбор примера

Модули и типы параметров у моделей различаются. На примере одной сети рассмотрим основной подход к преобразованию; для других моделей он будет аналогичным.

EfficientNet — работа Google, опубликованная в 2019 году. Подробная архитектура сети описана в статье. Здесь в качестве примера возьмём EfficientNet+FC модель с полносвязным слоем и рассмотрим преобразование параметров сети.

Загрузка модели 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 и получение параметров 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.

(Конец)

comments powered by Disqus