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)