时间戳

首页 > 解决方案 > 在 Pytorch 中为模型注册参数的正确方法

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

问题描述

我试图在 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()为模型实现梯度？

标签： pythonpytorch

解决方案

您注册参数过于复杂。您可以将一个新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)

推荐阅读