PyTorch设置与更新可训练参数

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)