将pytorch权重文件转化为mindspore权重文件的代码示例

将pytorch权重文件转化为mindspore权重文件的代码示例

1.取得参数名映射表

2.根据参数名映射迁移权重参数

代码示例如下：

# Copyright 2022 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Convert checkpoint from torch/facebook"""
import argparse
import mindspore as ms
import torch

def get_keymap_txt(pth_file):
    map_path = pth_file.split('.')[0] + '_key_map.txt'
    map_file = open(map_path, 'w')
    state_dict = torch.load(pth_file, map_location=torch.device('cpu'))
    if 'model_state' in state_dict:
        state_dict = state_dict['model_state']
    elif 'module' in state_dict:
        state_dict = state_dict['module']
    elif 'model' in state_dict:
        state_dict = state_dict['model']
    
    for k in state_dict:
        new_k = k
        if new_k.startswith('encoder.'):
            new_k = new_k.split('encoder.')[1]
        else:
            continue
        if new_k.startswith('norm.'):
            new_k = new_k.replace(