1、设置可训练参数

PyTorch中可以使用torch.nn.Parameter()来设置可训练参数。Parameter类是Tensor类的子类,当它与Module类一起使用时,也就是将一个Parameter对象作为Module类的一个属性时,它们会自动添加到Module的参数列表中,例如在该Module类的parameters()迭代器中。

import torch
import torch.nn as nn

class MyModule(nn.Module):
	def __init__(self, num):
		super(MyModule, self).__init__()
		params = torch.ones(num, requires_grad=True)
		self.params = nn.Parameter(params)
		
	def forward(self, x):
		y = self.params * x

my_module = MyModule(10)
inputs = torch.ones(10)
outputs = my_module(inputs)
print(my_module.state_dict())
print(list(my_module.parameters()))
print(dict(my_module.named_parameters()))

此时params自动加入了my_module的Parameters,打印结果:

OrderedDict([('params', tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]))])

[Parameter containing:
tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], requires_grad=True)]

{'params': Parameter containing:
tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], requires_grad=True)}

2、更新自定义的可训练参数

假如有一个网络net,训练时需要同时更新这个网络的参数和上面my_module的参数,可以这样定义优化器:

optimizer = optim.SGD([{"params": net.parameters()}, 
                       {"params": my_module.parameters()}], 
                        lr=base_lr, momentum=momentum, weight_decay=weight_decay)

 

Logo

华为开发者空间,是为全球开发者打造的专属开发空间,汇聚了华为优质开发资源及工具,致力于让每一位开发者拥有一台云主机,基于华为根生态开发、创新。

更多推荐