|
| 1 | +--- |
| 2 | +title: backward 操作时 grad_tensors 参数的作用 |
| 3 | +author: me |
| 4 | +date: 2024-10-25 20:45:00 +0800 |
| 5 | +categories: [Deep Learning, PyTorch] |
| 6 | +tags: [Deep Learning, PyTorch, Gradient] |
| 7 | +math: true |
| 8 | +--- |
| 9 | + |
| 10 | +## 问题提出 |
| 11 | + |
| 12 | +在阅读 Dive Into Deep Learning 时，我在 [ReLU激活函数部分](https://zh-v2.d2l.ai/chapter_multilayer-perceptrons/mlp.html#relu) 看到这样一段代码： |
| 13 | + |
| 14 | +```python |
| 15 | +import torch |
| 16 | +x = torch.arange(-8.0, 8.0, 0.1, requires_grad=True) |
| 17 | +y = torch.relu(x) |
| 18 | +y.backward(torch.ones_like(x), retain_graph=True) |
| 19 | +``` |
| 20 | + |
| 21 | +其中在对非标量张量 y 进行梯度计算时，引入了参数 `(grad_tensors=)torch.ones_like(x)`。如果不加入这个参数，在运行时则会报错： |
| 22 | + |
| 23 | +```python |
| 24 | +x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) |
| 25 | +y = x*x |
| 26 | +y.backward() |
| 27 | +``` |
| 28 | + |
| 29 | +```text |
| 30 | +RuntimeError: grad can be implicitly created only for scalar outputs |
| 31 | +``` |
| 32 | + |
| 33 | +证明在计算非标量张量对张量的梯度时，必须引入 `grad_tensors` 这一参数。但是为什么要引入这一参数？这个参数是如何参与到梯度计算中的？ |
| 34 | + |
| 35 | +阅读 [torch.autograd.backward 文档](https://pytorch.org/docs/stable/generated/torch.autograd.backward.html)，发现 PyTorch 这样解释计算非标量张量对张量的梯度的原理： |
| 36 | + |
| 37 | +>The graph is differentiated using the chain rule. If any of `tensors` are non-scalar (i.e. their data has more than one element) and require gradient, then the Jacobian-vector product would be computed, in this case the function additionally requires specifying `grad_tensors`. It should be a sequence of matching length, that contains the “vector” in the Jacobian-vector product, usually the gradient of the differentiated function w.r.t. corresponding tensors (`None` is an acceptable value for all tensors that don’t need gradient tensors). |
| 38 | +> |
| 39 | +>该（计算）图使用链式法则进行微分。如果任何`tensors`是非标量（即它们的数据具有多个元素）并且需要梯度，则将计算雅可比向量积，在这种情况下，该函数还需要指定`grad_tensors` 。它应该是一个匹配长度的序列，包含雅可比向量乘积中的“向量”，通常是对应张量的微分函数的梯度（对于所有不需要梯度张量的张量来说， `None`是可接受的值）。 |
| 40 | +
|
| 41 | +## 背景知识：雅可比矩阵 |
| 42 | + |
| 43 | +在一般的梯度计算中，如果要计算矩阵对矩阵的梯度，会借助雅各比矩阵来进行。例如，假设我们有一个函数 $$f(x)$$，其中： |
| 44 | +- 输入张量 $$x$$ 的形状为 $$[p]$$。 |
| 45 | +- 输出张量 $$y = f(x)$$ 的形状为 $$[n]$$。 |
| 46 | + |
| 47 | +那么，$$y$$ 对 $$x$$ 的雅可比矩阵 $$J$$ 的维度是 $$[n, p]$$。这表示 $$y$$ 中的每个元素 $$y_{i}$$ 对 $$x$$ 的偏导数 $$\frac{\partial y_{i}}{\partial x_k}$$ 构成了雅可比矩阵的元素 $$J_{i,k}$$。在这个雅可比矩阵中，张量 $$x$$ 的每个元素都对应了 $$n$$ 个梯度值。 |
| 48 | + |
| 49 | +$$ |
| 50 | +J = \begin{bmatrix} |
| 51 | +\frac{\partial y_1}{\partial x_1} & \frac{\partial y_1}{\partial x_2} & \cdots & \frac{\partial y_1}{\partial x_p} \\ |
| 52 | +\frac{\partial y_2}{\partial x_1} & \frac{\partial y_2}{\partial x_2} & \cdots & \frac{\partial y_2}{\partial x_p} \\ |
| 53 | +\vdots & \vdots & \ddots & \vdots \\ |
| 54 | +\frac{\partial y_n}{\partial x_1} & \frac{\partial y_n}{\partial x_2} & \cdots & \frac{\partial y_n}{\partial x_p} |
| 55 | +\end{bmatrix} |
| 56 | +$$ |
| 57 | + |
| 58 | +但这样做存在一些问题： |
| 59 | + |
| 60 | +1. 没有定义如何将每个 $$x_i$$ 的若干梯度值计算为单个梯度 |
| 61 | +2. 计算这么大的雅可比矩阵会消耗较大的内存和计算量。 |
| 62 | + |
| 63 | +所以在计算非标量张量对非标量张量的梯度时，PyTorch 不会直接计算雅可比矩阵（Jacobian Matrix），而是通过 `grad_tensors` 实现了对雅可比向量积的高效计算。这不仅节省内存和计算量，而且满足大多数深度学习应用的需求。 |
| 64 | + |
| 65 | +## 背景知识：雅可比向量积 |
| 66 | + |
| 67 | +**雅可比向量积**（JVP）则是将雅可比矩阵 $$J$$ 与一个向量 $$v$$ 相乘，即： |
| 68 | + |
| 69 | +$$ |
| 70 | +J \cdot v |
| 71 | +$$ |
| 72 | + |
| 73 | +其中 $$v$$ 的维度与 $$y$$ 相同。这一操作的结果是一个与 $$x$$ 相同维度的向量，即输出张量的加权导数之和。通过计算雅可比向量积，而不是完整的雅可比矩阵，我们可以避免计算和存储所有的偏导数。 |
| 74 | + |
| 75 | +$$ |
| 76 | +J \cdot v = \begin{pmatrix} |
| 77 | +\frac{\partial y_1}{\partial x_1} & \cdots & \frac{\partial y_n}{\partial x_1} \\ |
| 78 | +\vdots & \ddots & \vdots \\ |
| 79 | +\frac{\partial y_1}{\partial x_p} & \cdots & \frac{\partial y_n}{\partial x_p} |
| 80 | +\end{pmatrix} |
| 81 | +\begin{pmatrix} |
| 82 | +\frac{\partial l}{\partial y_1} \\ |
| 83 | +\vdots \\ |
| 84 | +\frac{\partial l}{\partial y_n} |
| 85 | +\end{pmatrix} |
| 86 | += \begin{pmatrix} |
| 87 | +\frac{\partial l}{\partial x_1} \\ |
| 88 | +\vdots \\ |
| 89 | +\frac{\partial l}{\partial x_p} |
| 90 | +\end{pmatrix} |
| 91 | +$$ |
| 92 | + |
| 93 | +## grad_tensors 和雅可比向量积的关系 |
| 94 | + |
| 95 | +在 PyTorch 中，`grad_tensors` 参数实际上就是充当这个向量 $$v$$ 的角色。在计算梯度时，如果我们传入了 `grad_tensors`，则 PyTorch 计算的梯度实际上是雅可比矩阵 $$J$$ 与 `grad_tensors` 的向量积 $$J \cdot \text{grad\_tensors}$$，也就是： |
| 96 | + |
| 97 | +$$ |
| 98 | +\frac{\partial (y \cdot \text{grad\_tensors}^{\mathrm T})}{\partial x} |
| 99 | +$$ |
| 100 | + |
| 101 | +这种方式高效地计算了雅可比向量积，而不是直接计算整个雅可比矩阵，从而避免了显式存储和计算高维的雅可比矩阵。 |
| 102 | + |
| 103 | +## 总结 |
| 104 | + |
| 105 | +通过使用 `grad_tensors`，PyTorch 可以高效地计算非标量张量对输入的梯度。其核心思想是利用雅可比向量积来避免显式构造雅可比矩阵，从而在内存和计算上实现了极大的优化。这种方式在深度学习中尤为有用，因为直接构造和存储高维雅可比矩阵不仅昂贵且不切实际。 |
| 106 | + |
| 107 | +## References |
| 108 | + |
| 109 | +1. [torch.autograd.backward - PyTorch Docs](https://pytorch.org/docs/stable/generated/torch.autograd.backward.html) |
| 110 | +2. [Pytorch autograd,backward详解](https://zhuanlan.zhihu.com/p/83172023) |
| 111 | +3. [我的 ChatGPT 提问过程](https://chatgpt.com/share/671b5b97-bd64-8007-a484-4394ed518190) |
0 commit comments