基于pytorch神经网络模型参数的加载及自定义

最近在训练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)