MxNet预训练模型到Pytorch模型的转换方式

 更新时间:2020年05月25日 14:21:32   作者:sparkexpert  
这篇文章主要介绍了MxNet预训练模型到Pytorch模型的转换方式,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧

预训练模型在不同深度学习框架中的转换是一种常见的任务。今天刚好DPN预训练模型转换问题,顺手将这个过程记录一下。

核心转换函数如下所示:

def convert_from_mxnet(model, checkpoint_prefix, debug=False):
 _, mxnet_weights, mxnet_aux = mxnet.model.load_checkpoint(checkpoint_prefix, 0)
 remapped_state = {}
 for state_key in model.state_dict().keys():
  k = state_key.split('.')
  aux = False
  mxnet_key = ''
  if k[0] == 'features':
   if k[1] == 'conv1_1':
    # input block
    mxnet_key += 'conv1_x_1__'
    if k[2] == 'bn':
     mxnet_key += 'relu-sp__bn_'
     aux, key_add = _convert_bn(k[3])
     mxnet_key += key_add
    else:
     assert k[3] == 'weight'
     mxnet_key += 'conv_' + k[3]
   elif k[1] == 'conv5_bn_ac':
    # bn + ac at end of features block
    mxnet_key += 'conv5_x_x__relu-sp__bn_'
    assert k[2] == 'bn'
    aux, key_add = _convert_bn(k[3])
    mxnet_key += key_add
   else:
    # middle blocks
    if model.b and 'c1x1_c' in k[2]:
     bc_block = True # b-variant split c-block special treatment
    else:
     bc_block = False
    ck = k[1].split('_')
    mxnet_key += ck[0] + '_x__' + ck[1] + '_'
    ck = k[2].split('_')
    mxnet_key += ck[0] + '-' + ck[1]
    if ck[1] == 'w' and len(ck) > 2:
     mxnet_key += '(s/2)' if ck[2] == 's2' else '(s/1)'
    mxnet_key += '__'
    if k[3] == 'bn':
     mxnet_key += 'bn_' if bc_block else 'bn__bn_'
     aux, key_add = _convert_bn(k[4])
     mxnet_key += key_add
    else:
     ki = 3 if bc_block else 4
     assert k[ki] == 'weight'
     mxnet_key += 'conv_' + k[ki]
  elif k[0] == 'classifier':
   if 'fc6-1k_weight' in mxnet_weights:
    mxnet_key += 'fc6-1k_'
   else:
    mxnet_key += 'fc6_'
   mxnet_key += k[1]
  else:
   assert False, 'Unexpected token'
 
  if debug:
   print(mxnet_key, '=> ', state_key, end=' ')
 
  mxnet_array = mxnet_aux[mxnet_key] if aux else mxnet_weights[mxnet_key]
  torch_tensor = torch.from_numpy(mxnet_array.asnumpy())
  if k[0] == 'classifier' and k[1] == 'weight':
   torch_tensor = torch_tensor.view(torch_tensor.size() + (1, 1))
  remapped_state[state_key] = torch_tensor
 
  if debug:
   print(list(torch_tensor.size()), torch_tensor.mean(), torch_tensor.std())
 
 model.load_state_dict(remapped_state)
 
 return model

从中可以看出,其转换步骤如下:

(1)创建pytorch的网络结构模型,设为model

(2)利用mxnet来读取其存储的预训练模型,得到mxnet_weights;

(3)遍历加载后模型mxnet_weights的state_dict().keys

(4)对一些指定的key值,需要进行相应的处理和转换

(5)对修改键名之后的key利用numpy之间的转换来实现加载。

为了实现上述转换,首先pip安装mxnet,现在新版的mxnet安装还是非常方便的。

第二步,运行转换程序,实现预训练模型的转换。

可以看到在相当的文件夹下已经出现了转换后的模型。

以上这篇MxNet预训练模型到Pytorch模型的转换方式就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持脚本之家。

相关文章

  • Python切片索引用法示例

    Python切片索引用法示例

    这篇文章主要介绍了Python切片索引用法,结合实例形式详细分析了Python切片索引的常见使用方法与操作注意事项,需要的朋友可以参考下
    2018-05-05
  • python 爬取吉首大学网站成绩单

    python 爬取吉首大学网站成绩单

    python简单爬虫,爬取吉首大学成绩单,学校的成绩单居然不支持导出,太坑了,算个绩点居然还要手打,我吐了.花2个多小时写了此python程序来生成可方便求和平均的Excel文件,帮助JSU学子脱离手算烦恼...
    2021-06-06
  • Pytorch配置GPU环境方式

    Pytorch配置GPU环境方式

    这篇文章主要介绍了Pytorch配置GPU环境方式,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教
    2024-02-02
  • python多进程并行代码实例

    python多进程并行代码实例

    这篇文章主要介绍了python多进程并行代码实例,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下
    2019-09-09
  • 分析解决Python中sqlalchemy数据库连接池QueuePool异常

    分析解决Python中sqlalchemy数据库连接池QueuePool异常

    这篇文章主要来给大家分析sqlalchemy数据库连接池QueuePool的异常,给大家用详细的图文方式做出了解决的方案,有需要的朋友可以借鉴参考下,希望可以有所帮助
    2021-09-09
  • Python之pandas读写文件乱码的解决方法

    Python之pandas读写文件乱码的解决方法

    下面小编就为大家分享一篇Python之pandas读写文件乱码的解决方法,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧
    2018-04-04
  • python 下载文件的几种方式分享

    python 下载文件的几种方式分享

    这篇文章主要介绍了python 下载文件的几种方式分享,帮助大家更好的理解和学习使用python,感兴趣的朋友可以了解下
    2021-04-04
  • python之Flask实现简单登录功能的示例代码

    python之Flask实现简单登录功能的示例代码

    这篇文章主要介绍了python之Flask实现简单登录功能的示例代码,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小编来一起学习学习吧
    2018-12-12
  • Selenium 安装和简单使用的实现

    Selenium 安装和简单使用的实现

    这篇文章主要介绍了Selenium 安装和简单使用的实现,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小编来一起学习学习吧
    2020-12-12
  • 如何运行带参数的python脚本

    如何运行带参数的python脚本

    这篇文章主要介绍了如何运行带参数的python脚本,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下
    2019-11-11

最新评论