PyTorch 2.2 中文官方教程（十一）（2）https://developer.aliyun.com/article/1482552
扩展 PyTorch
使用自定义函数进行双向传播
原文：
pytorch.org/tutorials/intermediate/custom_function_double_backward_tutorial.html译者：飞龙
有时候需要通过向后图两次运行反向传播，例如计算高阶梯度。然而，要支持双向传播需要对 autograd 有一定的理解和谨慎。支持单次向后传播的函数不一定能够支持双向传播。在本教程中，我们展示了如何编写支持双向传播的自定义 autograd 函数，并指出一些需要注意的事项。
当编写自定义 autograd 函数以进行两次向后传播时，重要的是要知道自定义函数中的操作何时被 autograd 记录，何时不被记录，以及最重要的是，save_for_backward 如何与所有这些操作一起使用。
自定义函数隐式影响梯度模式的两种方式：
- 在向前传播期间，autograd 不会记录任何在前向函数内执行的操作的图形。当前向完成时，自定义函数的向后函数将成为每个前向输出的 grad_fn
- 在向后传播期间，如果指定了 create_graph 参数，autograd 会记录用于计算向后传播的计算图
接下来，为了了解 save_for_backward 如何与上述交互，我们可以探索一些示例：
保存输入
考虑这个简单的平方函数。它保存一个输入张量以备向后传播使用。当 autograd 能够记录向后传播中的操作时，双向传播会自动工作，因此当我们保存一个输入以备向后传播时，通常不需要担心，因为如果输入是任何需要梯度的张量的函数，它应该有 grad_fn。这样可以正确传播梯度。
import torch class Square(torch.autograd.Function): @staticmethod def forward(ctx, x): # Because we are saving one of the inputs use `save_for_backward` # Save non-tensors and non-inputs/non-outputs directly on ctx ctx.save_for_backward(x) return x**2 @staticmethod def backward(ctx, grad_out): # A function support double backward automatically if autograd # is able to record the computations performed in backward x, = ctx.saved_tensors return grad_out * 2 * x # Use double precision because finite differencing method magnifies errors x = torch.rand(3, 3, requires_grad=True, dtype=torch.double) torch.autograd.gradcheck(Square.apply, x) # Use gradcheck to verify second-order derivatives torch.autograd.gradgradcheck(Square.apply, x)
我们可以使用 torchviz 来可视化图形以查看为什么这样可以工作
import torchviz x = torch.tensor(1., requires_grad=True).clone() out = Square.apply(x) grad_x, = torch.autograd.grad(out, x, create_graph=True) torchviz.make_dot((grad_x, x, out), {"grad_x": grad_x, "x": x, "out": out})
我们可以看到对于 x 的梯度本身是 x 的函数（dout/dx = 2x），并且这个函数的图形已经正确构建
保存输出
在前一个示例的轻微变化是保存输出而不是输入。机制类似，因为输出也与 grad_fn 相关联。
class Exp(torch.autograd.Function): # Simple case where everything goes well @staticmethod def forward(ctx, x): # This time we save the output result = torch.exp(x) # Note that we should use `save_for_backward` here when # the tensor saved is an ouptut (or an input). ctx.save_for_backward(result) return result @staticmethod def backward(ctx, grad_out): result, = ctx.saved_tensors return result * grad_out x = torch.tensor(1., requires_grad=True, dtype=torch.double).clone() # Validate our gradients using gradcheck torch.autograd.gradcheck(Exp.apply, x) torch.autograd.gradgradcheck(Exp.apply, x)
使用 torchviz 来可视化图形：
out = Exp.apply(x) grad_x, = torch.autograd.grad(out, x, create_graph=True) torchviz.make_dot((grad_x, x, out), {"grad_x": grad_x, "x": x, "out": out})
保存中间结果
更棘手的情况是当我们需要保存一个中间结果时。我们通过实现以下情况来演示这种情况：
s i n h ( x ) : = e x − e − x 2 sinh(x) := \frac{e^x - e^{-x}}{2}sinh(x):=2ex−e−x
由于 sinh 的导数是 cosh，因此在向后计算中重复使用 exp(x)和 exp(-x)这两个中间结果可能很有用。
尽管如此，中间结果不应直接保存并在向后传播中使用。因为前向是在无梯度模式下执行的，如果前向传递的中间结果用于计算向后传递中的梯度，则梯度的向后图将不包括计算中间结果的操作。这会导致梯度不正确。
class Sinh(torch.autograd.Function): @staticmethod def forward(ctx, x): expx = torch.exp(x) expnegx = torch.exp(-x) ctx.save_for_backward(expx, expnegx) # In order to be able to save the intermediate results, a trick is to # include them as our outputs, so that the backward graph is constructed return (expx - expnegx) / 2, expx, expnegx @staticmethod def backward(ctx, grad_out, _grad_out_exp, _grad_out_negexp): expx, expnegx = ctx.saved_tensors grad_input = grad_out * (expx + expnegx) / 2 # We cannot skip accumulating these even though we won't use the outputs # directly. They will be used later in the second backward. grad_input += _grad_out_exp * expx grad_input -= _grad_out_negexp * expnegx return grad_input def sinh(x): # Create a wrapper that only returns the first output return Sinh.apply(x)[0] x = torch.rand(3, 3, requires_grad=True, dtype=torch.double) torch.autograd.gradcheck(sinh, x) torch.autograd.gradgradcheck(sinh, x)
使用 torchviz 来可视化图形：
out = sinh(x) grad_x, = torch.autograd.grad(out.sum(), x, create_graph=True) torchviz.make_dot((grad_x, x, out), params={"grad_x": grad_x, "x": x, "out": out})
保存中间结果：不要这样做
现在我们展示当我们不返回中间结果作为输出时会发生什么：grad_x 甚至不会有一个反向图，因为它纯粹是一个函数 exp 和 expnegx，它们不需要 grad。
class SinhBad(torch.autograd.Function): # This is an example of what NOT to do! @staticmethod def forward(ctx, x): expx = torch.exp(x) expnegx = torch.exp(-x) ctx.expx = expx ctx.expnegx = expnegx return (expx - expnegx) / 2 @staticmethod def backward(ctx, grad_out): expx = ctx.expx expnegx = ctx.expnegx grad_input = grad_out * (expx + expnegx) / 2 return grad_input
使用 torchviz 来可视化图形。请注意，grad_x 不是图形的一部分！
out = SinhBad.apply(x) grad_x, = torch.autograd.grad(out.sum(), x, create_graph=True) torchviz.make_dot((grad_x, x, out), params={"grad_x": grad_x, "x": x, "out": out})
当不跟踪反向传播时
最后，让我们考虑一个例子，即 autograd 可能根本无法跟踪函数的反向梯度。我们可以想象 cube_backward 是一个可能需要非 PyTorch 库（如 SciPy 或 NumPy）或编写为 C++扩展的函数。这里演示的解决方法是创建另一个自定义函数 CubeBackward，在其中手动指定 cube_backward 的反向传播！
def cube_forward(x): return x**3 def cube_backward(grad_out, x): return grad_out * 3 * x**2 def cube_backward_backward(grad_out, sav_grad_out, x): return grad_out * sav_grad_out * 6 * x def cube_backward_backward_grad_out(grad_out, x): return grad_out * 3 * x**2 class Cube(torch.autograd.Function): @staticmethod def forward(ctx, x): ctx.save_for_backward(x) return cube_forward(x) @staticmethod def backward(ctx, grad_out): x, = ctx.saved_tensors return CubeBackward.apply(grad_out, x) class CubeBackward(torch.autograd.Function): @staticmethod def forward(ctx, grad_out, x): ctx.save_for_backward(x, grad_out) return cube_backward(grad_out, x) @staticmethod def backward(ctx, grad_out): x, sav_grad_out = ctx.saved_tensors dx = cube_backward_backward(grad_out, sav_grad_out, x) dgrad_out = cube_backward_backward_grad_out(grad_out, x) return dgrad_out, dx x = torch.tensor(2., requires_grad=True, dtype=torch.double) torch.autograd.gradcheck(Cube.apply, x) torch.autograd.gradgradcheck(Cube.apply, x)
使用 torchviz 来可视化图形：
out = Cube.apply(x) grad_x, = torch.autograd.grad(out, x, create_graph=True) torchviz.make_dot((grad_x, x, out), params={"grad_x": grad_x, "x": x, "out": out})
总之，双向传播是否适用于您的自定义函数取决于反向传播是否可以被 autograd 跟踪。通过前两个示例，我们展示了双向传播可以直接使用的情况。通过第三和第四个示例，我们展示了使反向函数可以被跟踪的技术，否则它们将无法被跟踪。
使用自定义函数融合卷积和批量归一化
原文：
pytorch.org/tutorials/intermediate/custom_function_conv_bn_tutorial.html译者：飞龙
注意
点击这里下载完整示例代码
将相邻的卷积和批量归一化层融合在一起通常是一种推理时间的优化，以提高运行时性能。通常通过完全消除批量归一化层并更新前面卷积的权重和偏置来实现[0]。然而，这种技术不适用于训练模型。
在本教程中，我们将展示一种不同的技术来融合这两个层，可以在训练期间应用。与改进运行时性能不同，这种优化的目标是减少内存使用。
这种优化的理念是看到卷积和批量归一化（以及许多其他操作）都需要在前向传播期间保存其输入的副本以供反向传播使用。对于大批量大小，这些保存的输入占用了大部分内存，因此能够避免为每个卷积批量归一化对分配另一个输入张量可以显著减少内存使用量。
在本教程中，我们通过将卷积和批量归一化合并为单个层（作为自定义函数）来避免这种额外的分配。在这个组合层的前向传播中，我们执行正常的卷积和批量归一化，唯一的区别是我们只保存卷积的输入。为了获得批量归一化的输入，这对于反向传播是必要的，我们在反向传播期间再次重新计算卷积的前向传播。
重要的是要注意，这种优化的使用是情境性的。虽然（通过避免保存一个缓冲区）我们总是在前向传播结束时减少分配的内存，但在某些情况下，峰值内存分配实际上可能并未减少。请查看最后一节以获取更多详细信息。
为简单起见，在本教程中，我们将 Conv2D 的 bias=False，stride=1，padding=0，dilation=1 和 groups=1 硬编码。对于 BatchNorm2D，我们将 eps=1e-3，momentum=0.1，affine=False 和 track_running_statistics=False 硬编码。另一个小的区别是在计算批量归一化时，在平方根的分母外部添加了 epsilon。
[0] nenadmarkus.com/p/fusing-batchnorm-and-conv/
卷积的反向传播公式实现
实现自定义函数需要我们自己实现反向传播。在这种情况下，我们需要为 Conv2D 和 BatchNorm2D 分别实现反向传播公式。最终，我们会将它们链接在一起形成统一的反向传播函数，但在下面，我们首先将它们实现为各自的自定义函数，以便验证它们的正确性
import torch from torch.autograd.function import once_differentiable import torch.nn.functional as F def convolution_backward(grad_out, X, weight): grad_input = F.conv2d(X.transpose(0, 1), grad_out.transpose(0, 1)).transpose(0, 1) grad_X = F.conv_transpose2d(grad_out, weight) return grad_X, grad_input class Conv2D(torch.autograd.Function): @staticmethod def forward(ctx, X, weight): ctx.save_for_backward(X, weight) return F.conv2d(X, weight) # Use @once_differentiable by default unless we intend to double backward @staticmethod @once_differentiable def backward(ctx, grad_out): X, weight = ctx.saved_tensors return convolution_backward(grad_out, X, weight)
在使用gradcheck进行测试时，重要的是使用双精度
weight = torch.rand(5, 3, 3, 3, requires_grad=True, dtype=torch.double) X = torch.rand(10, 3, 7, 7, requires_grad=True, dtype=torch.double) torch.autograd.gradcheck(Conv2D.apply, (X, weight))
True
PyTorch 2.2 中文官方教程（十一）（4）https://developer.aliyun.com/article/1482555#slide-3
