pytorch 在网络中添加可训练参数在网络中添加可训练参数,修改预训练权重文件的方法修改预训练权重文件的方法
今天小编就为大家分享一篇pytorch 在网络中添加可训练参数,修改预训练权重文件的方法,具有很好的参考价
值,希望对大家有所帮助。一起跟随小编过来看看吧
实践中,针对不同的任务需求,我们经常会在现成的网络结构上做一定的修改来实现特定的目的。
假如我们现在有一个简单的两层感知机网络:假如我们现在有一个简单的两层感知机网络:
# -*- coding: utf-8 -*-
import torch
from torch.autograd import Variable
import torch.optim as optim
x = Variable(torch.FloatTensor([1, 2, 3])).cuda()
y = Variable(torch.FloatTensor([4, 5])).cuda()
class MLP(torch.nn.Module):
def __init__(self):
super(MLP, self).__init__()
self.linear1 = torch.nn.Linear(3, 5)
self.relu = torch.nn.ReLU()
self.linear2 = torch.nn.Linear(5, 2)
def forward(self, x):
x = self.linear1(x)
x = self.relu(x)
x = self.linear2(x)
return x
model = MLP().cuda()
loss_fn = torch.nn.MSELoss(size_average=False)
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
for t in range(500):
y_pred = model(x)
loss = loss_fn(y_pred, y)
print(t, loss.data[0])
model.zero_grad()
loss.backward()
optimizer.step()
print(model(x))
现在想在前向传播时,在relu之后给x乘以一个可训练的系数,只需要在__init__函数中添加一个nn.Parameter类型变量,并在
forward函数中乘以该变量即可:
class MLP(torch.nn.Module):
def __init__(self):
super(MLP, self).__init__()
self.linear1 = torch.nn.Linear(3, 5)
self.relu = torch.nn.ReLU()
self.linear2 = torch.nn.Linear(5, 2)
# the para to be added and updated in train phase, note that NO cuda() at last
self.coefficient = torch.nn.Parameter(torch.Tensor([1.55]))
def forward(self, x):
x = self.linear1(x)
x = self.relu(x)
x = self.coefficient * x
x = self.linear2(x)
return x
注意,Parameter变量和Variable变量的操作大致相同,但是不能手动调用.cuda()方法将其加载在GPU上,事实上它会自动在
GPU上加载,可以通过model.state_dict()或者model.named_parameters()函数查看现在的全部可训练参数(包括通过继承得
到的父类中的参数):
print(model.state_dict().keys())
for i, j in model.named_parameters():
print(i)
print(j)
输出如下:
评论1