在Pytorch中为模型注册参数的正确方法

Question

我试图在Pytorch中定义一个简单的模型。该模型计算高斯分布的负对数问题：

import torch
import torch.nn as nn
class GaussianModel(nn.Module):
    def __init__(self):
        super(GaussianModel, self).__init__()
        self.register_parameter('mean', nn.Parameter(torch.zeros(1),
                                                     requires_grad=True))
        
        self.pdf = torch.distributions.Normal(self.state_dict()['mean'],
                                              torch.tensor([1.0]))
    def forward(self, x):
        return -self.pdf.log_prob(x)
model = GaussianModel()

然后我尝试优化mean参数：

optimizer = torch.optim.SGD(model.parameters(), lr=0.002)
for _ in range(5):
  optimizer.zero_grad()
  nll = model(torch.tensor([3.0], requires_grad=True))
  nll.backward()
  optimizer.step()
  print('mean : ', model.state_dict()['mean'],
                 ' - Negative Loglikelihood : ', nll.item())

但梯度似乎为零，mean没有变化：

mean :  tensor([0.])  - Negative Loglikelihood :  5.418938636779785
mean :  tensor([0.])  - Negative Loglikelihood :  5.418938636779785
mean :  tensor([0.])  - Negative Loglikelihood :  5.418938636779785
mean :  tensor([0.])  - Negative Loglikelihood :  5.418938636779785
mean :  tensor([0.])  - Negative Loglikelihood :  5.418938636779785

我是否正确注册并使用了mean参数？可以自动计算torch.distributions.Normal.log_prob的梯度吗？或者我应该为模型实现backward()？

Answer 1

您的参数注册过于复杂。您可以将一个新的self.mean属性指定为nn.Parameter，然后在大多数情况下将其用作张量。

nn.Module覆盖__setattr__方法，每次分配新的类属性时都会调用该方法。它所做的一件事是检查是否分配了nn.Parameter类型，如果是，它会将其添加到已注册参数的模块字典中。

因此，注册参数的最简单方法如下：

import torch
import torch.nn as nn
class GaussianModel(nn.Module):
    def __init__(self):
        super(GaussianModel, self).__init__()
        self.mean = nn.Parameter(torch.zeros(1))
        self.pdf = torch.distributions.Normal(self.mean, torch.tensor([1.0]))
    def forward(self, x):
        return -self.pdf.log_prob(x)