MxNet預訓練模型到Pytorch模型的轉(zhuǎn)換方式

更新時間：2020年05月25日 14:21:32 作者：sparkexpert

這篇文章主要介紹了MxNet預訓練模型到Pytorch模型的轉(zhuǎn)換方式，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧

預訓練模型在不同深度學習框架中的轉(zhuǎn)換是一種常見的任務。今天剛好DPN預訓練模型轉(zhuǎn)換問題，順手將這個過程記錄一下。

核心轉(zhuǎn)換函數(shù)如下所示：

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

從中可以看出，其轉(zhuǎn)換步驟如下：

（1）創(chuàng)建pytorch的網(wǎng)絡結構模型，設為model

（2）利用mxnet來讀取其存儲的預訓練模型，得到mxnet_weights；

（3）遍歷加載后模型mxnet_weights的state_dict().keys

（4）對一些指定的key值，需要進行相應的處理和轉(zhuǎn)換

（5）對修改鍵名之后的key利用numpy之間的轉(zhuǎn)換來實現(xiàn)加載。

為了實現(xiàn)上述轉(zhuǎn)換，首先pip安裝mxnet，現(xiàn)在新版的mxnet安裝還是非常方便的。