PyTorch：是否可以微分矩阵？

Question

如何在 PyTorch 中区分矩阵？我尝试过以下方法，但都不起作用：

实例1：

a = torch.tensor([1., 2, 3], requires_grad=True)
b = torch.tensor([4., 5, 6], requires_grad=True) 
c = a*b
c.backward()
#print(b.grad)

>>> RuntimeError: grad can be implicitly created only for scalar outputs

实例2：

a = torch.tensor([1., 2, 3], requires_grad=True)
b = torch.tensor([4., 5, 6], requires_grad=True)

c = a*b
print(b.grad)

>>> None

Answer 1

这是可能的，但它并不真正适合 PyTorch 的标准用例，您通常对标量值函数的梯度感兴趣。

矩阵 Y 对矩阵 X 的导数可以表示为广义雅可比行列式。对于两个矩阵都是向量的情况，这会简化为标准雅可比矩阵，其中雅可比矩阵的每一行都是 Y 的一个元素相对于 X 的梯度的转置。更一般地，如果 X 的形状为 (n1, n2, ..., nD) 并且 Y 的形状为 (m1, m2, ..., mE)，那么表示 Y 相对于 X 的广义雅可比行列式的自然方法是形状为 (m1, m2, ..., mE) 的张量, mE, n1, n2, ..., nD)。

据我所知，PyTorch 有两种计算广义雅可比行列式的方法。

选项1

对 Y 的每个元素重复应用反向传播。

import torch

def construct_jacobian(y, x, retain_graph=False):
    x_grads = []
    for idx, y_element in enumerate(y.flatten()):
        if x.grad is not None:
            x.grad.zero_()
        # if specified set retain_graph=False on last iteration to clean up
        y_element.backward(retain_graph=retain_graph or idx < y.numel() - 1)
        x_grads.append(x.grad.clone())
    return torch.stack(x_grads).reshape(*y.shape, *x.shape)

那么您的测试用例的雅可比行列式可以使用以下方式计算

a = torch.tensor([1., 2., 3.])
b = torch.tensor([4., 5., 6.], requires_grad=True)
c = a * b

jacobian = construct_jacobian(c, b)

print(jacobian)

这导致

tensor([[1., 0., 0.],
        [0., 2., 0.],
        [0., 0., 3.]])

选项2

在 PyTorch 1.5.1 中引入了新的 autograd.function API，包括新函数torch.autograd.functional.jacobian。这会产生与前面的示例相同的结果，但采用函数作为参数。此处未演示，但jacobian如果您的函数采用多个独立张量作为输入，则可以为函数提供输入列表。在这种情况下，jacobian将返回一个包含每个输入参数的广义雅可比行列式的元组。

import torch

a = torch.tensor([1., 2., 3.])

def my_fun(b):
    return a * b

b = torch.tensor([4., 5., 6.], requires_grad=True)

jacobian = torch.autograd.functional.jacobian(my_fun, b)

print(jacobian)

它还产生

tensor([[1., 0., 0.],
        [0., 2., 0.],
        [0., 0., 3.]])

顺便说一句，在一些文献中，术语“梯度”用于指雅可比矩阵的转置。如果这就是您想要的，假设 Y 和 X 是向量，您可以简单地使用上面的代码并对生成的雅可比矩阵进行转置。如果 Y 或 X 是高阶张量（矩阵或 n 维张量），那么我不知道有任何文献区分梯度和广义雅可比行列式。表示广义雅可比行列式的这种“转置”的一种自然方法是将Tensor.permute其转换为形状为 (n1, n2, ..., nD, m1, m2, ..., mE) 的张量。

另一方面，广义雅可比行列式的概念很少在文献中使用（示例用法），但实际上在实践中相对有用。这是因为它基本上作为一种簿记技术来跟踪 Y 和 X 的原始维度。我的意思是，您可以轻松地采用 Y 和 X 并将它们展平为向量，而不管它们的原始形状如何。那么导数将是标准雅可比矩阵。因此，该雅可比矩阵相当于广义雅可比矩阵的重塑版本。