【MindSpore】将PyTorch的checkpoint文件导入到MindSpore网络

原创

已于 2022-06-21 10:23:44 修改 · 2.3k 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#mindspore #pytorch

于 2021-05-08 11:55:47 首次发布

注：本文中用的是MindSpore1.3.0版本，参数的命名方式可能会有更新，请参照两框架的API文档自行甄别参数名的对应关系。（2022年6月21日）

我采用的是一种比较笨拙的方法，需要人工对比寻找规律，但是有效。

首先分别需要PyTorch和MindSpore的checkpoint文件，用下面的方式导出文件内容，遍历参数的key和value。

import torch
import mindspore 
from mindspore import load_checkpoint


def traversal_params(pth_file_path, ckpt_file_path):
    # load pth file as a dictionary
    torch_params_dict = torch.load(pth_file_path)
    # traversal a params dictionary
    for k, v in torch_params_dict.items():