PyTorch中对NN模型中的参数进行初始化

一、目前在尝试各种方法来训练自己的模型，希望能够得到较好的training_model，包括了对模型的初始化，记录一下代码部分：

from torch.nn import init
#define the initial function to init the layer's parameters for the network
def weigth_init(m):
   if isinstance(m, nn.Conv2d):
       init.xavier_uniform_(m.weight.data)
       init.constant_(m.bias.data,0.1)
   elif isinstance(m, nn.BatchNorm2d):
       m.weight.data.fill_(1)
       m.bias.data.zero_()
   elif isinstance(m, nn.Linear):
       m.weight.data.normal_(0,0.01)
       m.bias.data.zero_()

首先定义了一个初始化函数，接着进行调用就ok了，不过要先把网络模型实例化：

    #Define Network
    model = Net(args.input_channel,args.output_channel)
    model.apply(weigth_init)

此上就完成了对模型中训练参数的初始化。

在知乎上也有看到一个类似的版本，也相应的贴上来作为参考了：

def initNetParams(net):
    '''Init net parameters.'''
    for m in net.modules():
        if isinstance(m, nn.Conv2d):
            init.xavier_uniform(m.weight)
            if m.bias:
                init.constant(m.bias, 0)
        elif isinstance(m, nn.BatchNorm2d):
            init.constant(m.weight, 1)
            init.constant(m.bias, 0)
        elif isinstance(m, nn.Linear):
            init.normal(m.weight, std=1e-3)
            if m.bias:
                init.constant(m.bias, 0)
 
initNetParams(net)

二、模型的保存以及加载

1.保存有两种方式，第一种是保存模型的整个结构信息和参数，第二种是只保存模型的参数

1). torch.save(net, 'net.pkl')

2). torch.save(net.state_dict(), 'net_params.pkl')

2.加载对应保存的两种网络

# 保存和加载整个模型  
torch.save(model_object, 'model.pth')  
model = torch.load('model.pth')  
 
# 仅保存和加载模型参数  
torch.save(model_object.state_dict(), 'params.pth')  
model_object.load_state_dict(torch.load('params.pth'))  

参考原文：https://blog.csdn.net/remanented/article/details/89161297