pytorch保存、加载参数的方式,初始化权重方式

假设网络为model = Net(), optimizer = optim.Adam(model.parameters(), lr=args.lr), 假设在某个epoch，我们要保存模型参数，优化器参数以及epoch

一、参数保存方式

1. 先建立一个字典，保存三个参数：

state = {‘net':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}

2.调用torch.save():

torch.save(state, dir)

其中dir表示保存文件的绝对路径+保存文件名，如'/home/qinying/Desktop/modelpara.pth'

或者这样写

torch.save({

'epoch': epoch + 1,

扫描二维码关注公众号，回复： 6211586 查看本文章

'arch': args.arch,

'state_dict': model.state_dict(),

'best_prec1': best_prec1,

}, 'checkpoint.tar' )

二、参数加载方式

当你想恢复某一阶段的训练（或者进行测试）时，那么就可以读取之前保存的网络模型参数等。

checkpoint = torch.load(dir)

model.load_state_dict(checkpoint['net'])

optimizer.load_state_dict(checkpoint['optimizer'])

start_epoch = checkpoint['epoch'] + 1

三、初始化权重方法

1. 均匀分布
torch.nn.init.uniform_(tensor, a=0, b=1)
服从~U(a,b) U(a, b)U(a,b)

2. 正太分布
torch.nn.init.normal_(tensor, mean=0, std=1)
服从~N(mean,std) N(mean, std)N(mean,std)

3. 初始化为常数
torch.nn.init.constant_(tensor, val)
初始化整个矩阵为常数val

4. Xavier
基本思想是通过网络层时，输入和输出的方差相同，包括前向传播和后向传播。

对于Xavier初始化方式，pytorch提供了uniform和normal两种方式

具体实例：

def weights_init(m):
    classname=m.__class__.__name__
    if classname.find('Conv') != -1:
        xavier(m.weight.data)
        xavier(m.bias.data)
net = Net()#构建网络
net.apply(weights_init) #apply函数会递归地搜索网络内的所有module并把参数表示的函数应用到所有的module上。  
 #对所有的Conv层都初始化权重.

import torch.nn.init as init
def weights_init(m):
    if type(m) in [nn.ConvTranspose2d, nn.Conv2d]:
        init.xavier_normal_(m.weight)
    elif type(m) == nn.BatchNorm2d:
        init.normal_(m.weight, 1.0, 0.02)
        init.constant_(m.bias, 0)
        
gnet.apply(weights_init)
dnet.apply(weights_init)

for m in model.modules():
    if isinstance(m,nn.Conv2d):
        nn.init.normal(m.weight.data)
        nn.init.xavier_normal(m.weight.data)
        nn.init.kaiming_normal(m.weight.data)#卷积层参数初始化
        m.bias.data.fill_(0)
    elif isinstance(m,nn.Linear):
        m.weight.data.normal_()#全连接层参数初始化

pytorch保存、加载参数的方式,初始化权重方式

猜你喜欢