python - 如何使用 pytorch 获取 jacobian 以获得多元正态分布的对数概率

Question

我已经从多元正态分布中抽取样本，并希望获得它们的对数概率相对于均值的梯度。由于有很多样本，这需要一个雅可比行列式：

import torch

mu = torch.ones((2,), requires_grad=True)
sigma = torch.eye(2)
dist = torch.distributions.multivariate_normal.MultivariateNormal(mu, sigma)

num_samples=10
samples = dist.sample((num_samples,))
logprobs = dist.log_prob(samples)

现在我想获得每个条目logprobs相对于mu.

一个简单的解决方案是 python 循环：

grads = []
for logprob in logprobs:
    grad = torch.autograd.grad(logprob, mu, retain_graph=True)
    grads.append(grad)

如果你堆叠毕业生，结果是所需的雅可比行列式。是否还有对此的内置和矢量化支持？

相关问题/互联网资源：

这是一个很大的话题，有很多相关的帖子。尽管如此，我认为这个具体问题（关于分布）尚未得到回答：

• 这个问题与我的基本相同（但没有示例代码和解决方案尝试），遗憾的是没有答案：Pytorch custom function jacobian gradient

• This question显示了pytorch中雅可比的计算，但我认为该解决方案不适用于我的问题：Pytorch最有效的雅可比/黑森计算它需要以一种似乎与分布不兼容的方式堆叠输入。我无法让它工作。

• 这个要点有一些雅可比的代码片段。原则上，它们类似于上述问题的方法。

标签： pythonpytorch

Answer 1

PyTorch 1.5.1 引入了一个torch.autograd.functional.jacobian函数。这将计算输入张量的函数的雅可比行列式。由于jacobian需要一个 python 函数作为第一个参数，因此使用它需要一些代码重组。

import torch

torch.manual_seed(0)    # for repeatable results

mu = torch.ones((2,), requires_grad=True)
sigma = torch.eye(2)
num_samples = 10

def f(mu):
    dist = torch.distributions.multivariate_normal.MultivariateNormal(mu, sigma)
    samples = dist.sample((num_samples,))
    logprobs = dist.log_prob(samples)
    return logprobs

grads = torch.autograd.functional.jacobian(f, mu)

print(grads)

tensor([[-1.1258, -1.1524],
        [-0.2506, -0.4339],
        [ 0.5988, -1.5551],
        [-0.3414,  1.8530],
        [ 0.4681, -0.1577],
        [ 1.4437,  0.2660],
        [ 1.3894,  1.5863],
        [ 0.9463, -0.8437],
        [ 0.9318,  1.2590],
        [ 2.0050,  0.0537]])