最近在训练MobileNet时经常会对其模型参数进行各种操作,或者替换其中的几层之类的,故总结一下用到的对神经网络参数的各种操作方法。
1.将matlab的.mat格式参数整理转换为tensor类型的模型参数
import torch import torch.nn as nn import torch.nn.functional as F import scipy.io as scio import os import numpy as np from config import Config import json config = Config() Mul = Config.MUL.astype('float32') Shift = Config.SHIFT.astype('float32') def load_json(j_fn): with open(j_fn,'r') as f: data = json.load(f) return data def save_json(dic,j_fn): json_str = json.dumps(dic) with open(j_fn,'w') as json_file: json_file.write(json_str) w_dic = {} b_dic = {} dic_all = {} for i in range(1,28,2): a = 'w'+str(i) #按顺序命名 b = 'b'+str(i) dic_all[a] = torch.from_numpy(scio.loadmat(config.WEIGHT_PATH + str(i)+'.mat')['wei'] * Mul[i-1]/(2**Shift[i-1])).permute(3, 2, 0, 1) dic_all[b] = torch.squeeze(torch.from_numpy(scio.loadmat(config.BIAS_PATH + str(i)+'.mat')['bias'] * Mul[i-1]/(2**Shift[i-1]))) # print(a, 'Mul'+str(i-1)) if i == 27: break a = 'w'+str(i+1) b = 'b'+str(i+1) dic_all[a] = torch.from_numpy(scio.loadmat(config.WEIGHT_PATH + str(i+1)+'.mat')['wei'] * Mul[i]/(2**Shift[i])).permute(2, 0, 1).unsqueeze(1) dic_all[b] = torch.squeeze(torch.from_numpy(scio.loadmat(config.BIAS_PATH + str(i+1)+'.mat')['bias'] * Mul[i]/(2**Shift[i]))) #此处由于自己之前的命名问题,中间跳过了28层(池化层),直接按照有参数的层存储了参数,故27后的文件名变成了29 dic_all['w29'] = torch.squeeze(torch.from_numpy(scio.loadmat(config.WEIGHT_PATH + '29.mat')['wei'] * Mul[28]/(2**Shift[28])).permute(3, 2, 0, 1)[1:, :]) dic_all['b29'] = torch.squeeze(torch.from_numpy(scio.loadmat(config.BIAS_PATH + '29.mat')['bias'] * Mul[28]/(2**Shift[28])))[1:] #存为.pth文件 param_fn = 'mobilenet_param_float.pth' torch.save(dic_all,param_fn)
其中,mul和shift为量化后的乘子和移位参数(如果参数是浮点的则可以忽略这部分),另外,我的weight和bias是按照每层单独存在一个按照层序号命名的.mat文件中。且由于是从matlab的程序得到的,需要对参数的维度进行一下转换(permute()方法),同时对需要增加或减少维度的用unsqueeze()或torch.squeeze()方法进行改变(注意一定要和网络需要的输入维度相同才行)。最后按照原来对参数文件命名的方式保存成字典存成.pth文件(此时的字典还不能直接使用,需要在具体定义的网络中更换想应的key值)。
*另外,代码中用来读取和存储.json文件的函数可以忽略,在这里没有用到
2.将自定义网络的参数替换成自己需要的(DIY模型参数)
from mobilenet_v1 import MobileNet_v1 import torch from config import Config from load_data import loadtestdata from torch.autograd import Variable import numpy as np from Mobilenetv1_quantified import MobileNet, MobileNet_Bayer import json import matplotlib.pyplot as plt import numpy as np import torchvision param_keys = ['w1', 'b1', 'w2', 'b2', 'w3', 'b3', 'w4', 'b4', 'w5', 'b5', 'w6', 'b6', 'w7', 'b7', 'w8', 'b8', 'w9', 'b9', 'w10', 'b10', 'w11', 'b11', 'w12', 'b12', 'w13', 'b13', 'w14', 'b14', 'w15', 'b15', 'w16', 'b16', 'w17', 'b17', 'w18', 'b18', 'w19', 'b19', 'w20', 'b20', 'w21', 'b21', 'w22', 'b22', 'w23', 'b23', 'w24', 'b24', 'w25', 'b25', 'w26', 'b26', 'w27', 'b27', 'w29', 'b29'] file_name = '/home/wangshuyu/MobileNet_v1/mobilenet_param_float.pth' dic_param = torch.load(file_name) # 此处打开上一步存成的参数字典(按照每一层的权重、偏置的顺序) Model = MobileNet() # 实例化预定义的MobileNet网络(网络结构将在其他文中给出) net_dic = Model.state_dict() # 加载预定义网络的参数字典,用来获取网络的键值 for i, param_tensor in enumerate(net_dic ,0): net_dic[param_tensor] = dic_param[param_keys[i]] # print(i,'\t',param_tensor ,net_dic[param_tensor].shape) #可以用来查看参数的维度 param_fn = 'MobileNet_float.pth' torch.save(net_dic,param_fn) # 下面开始是自己定义的另一个网络,只需要固定MobileNet其中一部分参数,剩下的部分参数用来训练,因此只从第11个之后的开始取参数 model2 = MobileNet_Bayer() dic2 = model2.state_dict() key_list = list(net_dic.keys()) for i, param_tensor in enumerate(dic2 ,0): if i > 11: dic2[param_tensor] = (net_dic[key_list[i - 2]]) print(i, '\t', param_tensor, dic2[param_tensor].shape) param_fn2 = 'MobileNet_Bayer.pth' torch.save(dic2,param_fn2)
这里主要实现了将之前存好的量化后的mobilenet每层参数根据自己定义的网络构建了参数字典,在训练或测试的时候,只需要加载之前存好的预训练参数就可以了:
from Mobilenetv1_quantified import MobileNet import torch
from load_data import loadtestdata Net = MobileNet() param_dic = torch.load('MobileNet_float.pth') Net.load_state_dict(param_dic)
classes = range(0,1000)
test_data = loadtestdate()
test(test_data, Net, classes)
本站文章如无特殊说明,均为本站原创,如若转载,请注明出处:基于pytorch神经网络模型参数的加载及自定义 - Python技术站