Tensorflow的MobileNetV1参数迁移到pytorch上并保存
因为放弃tensorflow超级久了，也不想再去用它，因为明明很简单用pytorch十几行作出的代码，tensorflow的版本完全看不懂，我这个菜鸡还是老老实实刨地吧。mobilenet的代码网上一大堆，我把我写的贴出来吧，论文简单易读，连我这种英语渣渣两天就看完了。
mobelnet的代码如下。
import torch.nn as nn import torch class Conv_bn(nn.Module): def __init__(self,inp,oup,stride): super(Conv_bn, self).__init__() self.convBn=nn.Sequential( nn.Conv2d(inp,oup,3,stride,1,bias=False), nn.BatchNorm2d(oup), nn.ReLU(inplace=True) ) def forward(self,x): out=self.convBn(x) return out class Conv_depth(nn.Module): def __init__(self,inp,oup,stride): super(Conv_depth, self).__init__() self.convDepthwise=nn.Sequential( nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False), nn.BatchNorm2d(inp), nn.ReLU(inplace=True), nn.Conv2d(inp, oup, 1, 1, 0, bias=False), nn.BatchNorm2d(oup), nn.ReLU(inplace=True) ) def forward(self,x): out=self.convDepthwise(x) return out class MobileNet(nn.Module): def __init__(self): super(MobileNet, self).__init__() self.mobelnet=nn.Sequential( Conv_bn(3, 32, 2), Conv_depth(32, 64, 1), Conv_depth(64, 128, 2), Conv_depth(128, 128, 1), Conv_depth(128, 256, 2), Conv_depth(256, 256, 1), Conv_depth(256, 512, 2), Conv_depth(512, 512, 1), Conv_depth(512, 512, 1), Conv_depth(512, 512, 1), Conv_depth(512, 512, 1), Conv_depth(512, 512, 1), Conv_depth(512, 1024, 2), Conv_depth(1024, 1024, 1), nn.AvgPool2d(7),) self.fc = nn.Linear(1024, 1000) # 网络的前向过程 def forward(self, x): x=self.mobelnet(x) x=x.view(-1, 1024) x=self.fc(x) return x
妈呀，简单吧，但是你不知道tensorflow的版本有多长啊。
然后转参数把我难住了，没做过，参考了 https://www.jianshu.com/p/0a61caeb693b 这位同学的moielnetV3版本的改法，但是我真的不懂他那个字典怎么定义的，我每次model.层名 就开始给我出红杠杠，报错，我估计可能是他把层都封装成了对象吧，如果有懂的同学希望能给我讲讲哈。我贴我自己的代码吧。
import json import tensorflow as tf import os from MobileNet.mobilenet_v1 import MobileNet import numpy as np import torch import os os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' CHECKPOINT_PATH='/Users/wenyu/Desktop/TorchProject/MobileNet/mobilenet_v1_1.0_224/mobilenet_v1_1.0_224.ckpt' # write the json file def new_dict(checkpoint_path,json_path): reader=tf.compat.v1.train.NewCheckpointReader(checkpoint_path) weights_shape =reader.get_variable_to_shape_map() print('the layer',weights_shape['MobilenetV1/Conv2d_9_pointwise/BatchNorm/moving_mean']) length=len(weights_shape['MobilenetV1/Conv2d_9_pointwise/BatchNorm/moving_mean']) # print(length) if not os.path.exists(json_path): weights_small = {n: 1 for (n, _) in reader.get_variable_to_shape_map().items()} keys_list=list(weights_small.keys()) for key_ in keys_list: if "/ExponentialMovingAverage" in key_: del weights_small[key_] elif "/RMSProp" in key_: del weights_small[key_] with open(json_path, 'w') as writer: json.dump(weights_small, fp=writer, sort_keys=True) else: print('the json file has been write!') # get convBn_dict def get_convbn_convert_dict(layer_num): convert_dict={ 'mobelnet.'+str(layer_num)+'.convBn.0.weight':'MobilenetV1/Conv2d_'+str(layer_num)+'/weights', 'mobelnet.'+str(layer_num)+'.convBn.1.weight':'MobilenetV1/Conv2d_'+str(layer_num)+'/BatchNorm/beta', 'mobelnet.'+str(layer_num)+'.convBn.1.bias':'MobilenetV1/Conv2d_'+str(layer_num)+'/BatchNorm/gamma', 'mobelnet.'+str(layer_num)+'.convBn.1.running_mean':'MobilenetV1/Conv2d_'+str(layer_num)+'/BatchNorm/moving_mean', 'mobelnet.'+str(layer_num)+'.convBn.1.running_var':'MobilenetV1/Conv2d_'+str(layer_num)+'/BatchNorm/moving_variance' } return convert_dict # get depthWise_dict def get_dpwise_convert_dict(layer_num): convert_dict={ 'mobelnet.'+str(layer_num)+'.convDepthwise.0.weight': 'MobilenetV1/Conv2d_'+str(layer_num)+'_depthwise/depthwise_weights', 'mobelnet.'+str(layer_num)+'.convDepthwise.1.weight': 'MobilenetV1/Conv2d_'+str(layer_num)+'_depthwise/BatchNorm/beta', 'mobelnet.'+str(layer_num)+'.convDepthwise.1.bias': 'MobilenetV1/Conv2d_'+str(layer_num)+'_depthwise/BatchNorm/gamma', 'mobelnet.'+str(layer_num)+'.convDepthwise.1.running_mean': 'MobilenetV1/Conv2d_'+str(layer_num)+'_depthwise/BatchNorm/moving_mean', 'mobelnet.'+str(layer_num)+'.convDepthwise.1.running_var': 'MobilenetV1/Conv2d_'+str(layer_num)+'_depthwise/BatchNorm/moving_variance', 'mobelnet.'+str(layer_num)+'.convDepthwise.3.weight': 'MobilenetV1/Conv2d_'+str(layer_num)+'_pointwise/weights', 'mobelnet.'+str(layer_num)+'.convDepthwise.4.weight': 'MobilenetV1/Conv2d_'+str(layer_num)+'_pointwise/BatchNorm/beta', 'mobelnet.' + str(layer_num) + '.convDepthwise.4.bias': 'MobilenetV1/Conv2d_' + str(layer_num) + '_pointwise/BatchNorm/gamma', 'mobelnet.' + str(layer_num) + '.convDepthwise.4.running_mean': 'MobilenetV1/Conv2d_' + str(layer_num) + '_pointwise/BatchNorm/moving_mean', 'mobelnet.' + str(layer_num) + '.convDepthwise.4.running_var': 'MobilenetV1/Conv2d_' + str(layer_num) + '_pointwise/BatchNorm/moving_variance' } return convert_dict # get conversion_dict def get_model_dict(layers_num): merge = lambda dict1, dict2: {**dict1, **dict2} conversion_table = {} convBn_dict=get_convbn_convert_dict(0) conversion_table=merge(conversion_table,convBn_dict) for i in range(1,layers_num): dpWise_dict=get_dpwise_convert_dict(i) conversion_table=merge(conversion_table,dpWise_dict) # load_parameter(CHECKPOINT_PATH,conversion_table) return conversion_table def write_json(conversion_table,json_path): if not os.path.exists(json_path): with open(json_path, 'w') as writer: json.dump(conversion_table, fp=writer, sort_keys=True) else: print('the conversion table has been wirten!') def load_parameter(conversion_table): module=MobileNet() original_model_dict=module.state_dict() pth_list=list(conversion_table.keys()) ckpt_list=list(conversion_table.values()) assert len(pth_list)==len(ckpt_list) ,('the length is not right!') reader=tf.compat.v1.train.NewCheckpointReader(CHECKPOINT_PATH) for i,ckpt_name in enumerate(ckpt_list): ckpt_name_value=tf.compat.v1.train.load_variable(CHECKPOINT_PATH,ckpt_name) if 'Conv2d' in ckpt_name and 'weights' in ckpt_name: ckpt_name_value=np.transpose(ckpt_name_value,(3,2,0,1)) if 'depthwise' in ckpt_name: ckpt_name_value=np.transpose(ckpt_name_value,(1,0,2,3)) elif 'BatchNorm' in ckpt_name and ckpt_name_value.ndim==1: # ckpt_name_value=np.transpose(ckpt_name_value) ckpt_name_value=ckpt_name_value pytorch_dict_key=pth_list[i] original_model_dict[pytorch_dict_key].data=torch.from_numpy(ckpt_name_value) torch.save(original_model_dict,'/Users/wenyu/Desktop/TorchProject/MobileNet/tf_to_torch.pth') return original_model_dict if __name__ == '__main__': conversion_table=get_model_dict(14) dic_mobel=load_parameter(conversion_table) print(dic_mobel['mobelnet.1.convDepthwise.0.weight'].shape)
其中核心就在最后两个函数，可能代码看起来很简单，但是我想了好久要怎么做，第一次做很不熟练，但是通过这次巩固了很多numpy，tensor还有字典的基本知识，很充实。有问题可以在博客下面留言。
