pytorch 在网络中添加可训练参数,修改预训练权重文件的方法
实践中,针对不同的任务需求,我们经常会在现成的网络结构上做一定的修改来实现特定的目的。
假如我们现在有一个简单的两层感知机网络:
#-*-coding:utf-8-*- importtorch fromtorch.autogradimportVariable importtorch.optimasoptim x=Variable(torch.FloatTensor([1,2,3])).cuda() y=Variable(torch.FloatTensor([4,5])).cuda() classMLP(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) defforward(self,x): x=self.linear1(x) x=self.relu(x) x=self.linear2(x) returnx model=MLP().cuda() loss_fn=torch.nn.MSELoss(size_average=False) optimizer=optim.SGD(model.parameters(),lr=0.001,momentum=0.9) fortinrange(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函数中乘以该变量即可:
classMLP(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) #theparatobeaddedandupdatedintrainphase,notethatNOcuda()atlast self.coefficient=torch.nn.Parameter(torch.Tensor([1.55])) defforward(self,x): x=self.linear1(x) x=self.relu(x) x=self.coefficient*x x=self.linear2(x) returnx
注意,Parameter变量和Variable变量的操作大致相同,但是不能手动调用.cuda()方法将其加载在GPU上,事实上它会自动在GPU上加载,可以通过model.state_dict()或者model.named_parameters()函数查看现在的全部可训练参数(包括通过继承得到的父类中的参数):
print(model.state_dict().keys()) fori,jinmodel.named_parameters(): print(i) print(j)
输出如下:
odict_keys(['linear1.weight','linear1.bias','linear2.weight','linear2.bias']) linear1.weight Parametercontaining: -0.3582-0.02830.2607 0.5190-0.22210.0665 -0.2586-0.33110.1927 -0.27650.5590-0.2598 0.4679-0.2923-0.3379 [torch.cuda.FloatTensorofsize5x3(GPU0)] linear1.bias Parametercontaining: -0.2549 -0.5246 -0.1109 0.5237 -0.1362 [torch.cuda.FloatTensorofsize5(GPU0)] linear2.weight Parametercontaining: -0.0286-0.30450.1928-0.23230.2966 0.26010.1441-0.21590.24840.0544 [torch.cuda.FloatTensorofsize2x5(GPU0)] linear2.bias Parametercontaining: -0.4038 0.3129 [torch.cuda.FloatTensorofsize2(GPU0)]
这个参数会在反向传播时与原有变量同时参与更新,这就达到了添加可训练参数的目的。
如果我们有原先网络的预训练权重,现在添加了一个新的参数,原有的权重文件自然就不能加载了,我们需要修改原权重文件,在其中添加我们的新变量的初始值。
调用model.state_dict查看我们添加的参数在参数字典中的完整名称,然后打开原先的权重文件:
a=torch.load("OldWeights.pth")a是一个collecitons.OrderedDict类型变量,也就是一个有序字典,直接将新参数名称和初始值作为键值对插入,然后保存即可。
a=torch.load("OldWeights.pth") a["layer1.0.coefficient"]=torch.FloatTensor([1.2]) a["layer1.1.coefficient"]=torch.FloatTensor([1.5]) torch.save(a,"Weights.pth")
现在权重就可以加载在修改后的模型上了。
以上这篇pytorch在网络中添加可训练参数,修改预训练权重文件的方法就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持毛票票。