mxnet模型转pytorch

最新推荐文章于 2025-06-01 09:15:13 发布

原创

最新推荐文章于 2025-06-01 09:15:13 发布 · 1.3k 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#mxnet #pytorch

本文详细介绍了从MXNet模型转换到PyTorch模型的过程，包括模型的读取、键名匹配、键值复制及测试保存。特别关注了BN层在两个框架之间的差异及其解决方法。

文章目录

- 核心:模型序列字典中键名的匹配

核心:模型序列字典中键名的匹配

1 读取mxnet模型

import mxnet as mx
import torch
from resnet50_gcn import resnet50
from torch.nn import BatchNorm2d
from torch.nn import Conv2d
from torch.nn import Linear
import pandas as pd
import numpy as np

# 加载符号图与模型参数
def get_model(model_path, epoch):
    sym, arg_params, aux_params = mx.model.load_checkpoint(model_path, epoch)
    return sym, arg_params,aux_params

# read mxnet model
sym, arg_params, aux_params = get_model('resnet50_w-glore_0-3_', 0)

2 读取pytorch模型

pytorch_model = resnet50(pretrained=False)
pytorch_model = init_model(pytorch_model, new_arg_params)

3 键名匹配

def dict_change_keys(input_dict, change_map, new_dict):
    # new_dict = {}
    old_dict_key = []
    counts_num = 0
    # old_dict_key = list(pd.read_csv(change_map, header=None, index_col=False, sep='')[0])
    new_dict_key = pd.read_csv(change_map, header=None, sep=' ')
    for key in