在 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参数？autograd 可以计算梯度吗？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)