PyTorch 2.2 中文官方教程（十）（3）https://developer.aliyun.com/article/1482541
计算 Hessian 向量积
计算 Hessian 向量积的朴素方法是将完整的 Hessian 材料化并与向量进行点积。我们可以做得更好：事实证明，我们不需要材料化完整的 Hessian 来做到这一点。我们将介绍两种（许多种）不同的策略来计算 Hessian 向量积：-将反向模式 AD 与反向模式 AD 组合-将反向模式 AD 与正向模式 AD 组合
将反向模式 AD 与正向模式 AD 组合（而不是反向模式与反向模式）通常是计算 HVP 的更节省内存的方式，因为正向模式 AD 不需要构建 Autograd 图并保存反向传播的中间结果：
from torch.func import jvp, grad, vjp def hvp(f, primals, tangents): return jvp(grad(f), primals, tangents)[1]
以下是一些示例用法。
def f(x): return x.sin().sum() x = torch.randn(2048) tangent = torch.randn(2048) result = hvp(f, (x,), (tangent,))
如果 PyTorch 正向 AD 没有覆盖您的操作，那么我们可以将反向模式 AD 与反向模式 AD 组合：
def hvp_revrev(f, primals, tangents): _, vjp_fn = vjp(grad(f), *primals) return vjp_fn(*tangents) result_hvp_revrev = hvp_revrev(f, (x,), (tangent,)) assert torch.allclose(result, result_hvp_revrev[0])
脚本的总运行时间：（0 分钟 10.644 秒）
下载 Python 源代码：jacobians_hessians.py
下载 Jupyter 笔记本：jacobians_hessians.ipynb
模型集成
原文：
pytorch.org/tutorials/intermediate/ensembling.html译者：飞龙
注意
点击这里下载完整的示例代码
这个教程演示了如何使用torch.vmap来对模型集合进行向量化。
什么是模型集成？
模型集成将多个模型的预测组合在一起。传统上，这是通过分别在一些输入上运行每个模型，然后组合预测来完成的。然而，如果您正在运行具有相同架构的模型，则可能可以使用torch.vmap将它们组合在一起。vmap是一个函数变换，它将函数映射到输入张量的维度。它的一个用例是通过向量化消除 for 循环并加速它们。
让我们演示如何使用简单 MLP 的集成来做到这一点。
注意
这个教程需要 PyTorch 2.0.0 或更高版本。
import torch import torch.nn as nn import torch.nn.functional as F torch.manual_seed(0) # Here's a simple MLP class SimpleMLP(nn.Module): def __init__(self): super(SimpleMLP, self).__init__() self.fc1 = nn.Linear(784, 128) self.fc2 = nn.Linear(128, 128) self.fc3 = nn.Linear(128, 10) def forward(self, x): x = x.flatten(1) x = self.fc1(x) x = F.relu(x) x = self.fc2(x) x = F.relu(x) x = self.fc3(x) return x
让我们生成一批虚拟数据，并假装我们正在处理一个 MNIST 数据集。因此，虚拟图像是 28x28，我们有一个大小为 64 的小批量。此外，假设我们想要将来自 10 个不同模型的预测组合起来。
device = 'cuda' num_models = 10 data = torch.randn(100, 64, 1, 28, 28, device=device) targets = torch.randint(10, (6400,), device=device) models = [SimpleMLP().to(device) for _ in range(num_models)]
我们有几种选项来生成预测。也许我们想给每个模型一个不同的随机小批量数据。或者，也许我们想通过每个模型运行相同的小批量数据（例如，如果我们正在测试不同模型初始化的效果）。
选项 1：为每个模型使用不同的小批量
minibatches = data[:num_models] predictions_diff_minibatch_loop = [model(minibatch) for model, minibatch in zip(models, minibatches)]
选项 2：相同的小批量
minibatch = data[0] predictions2 = [model(minibatch) for model in models]
使用vmap来对集合进行向量化
让我们使用vmap来加速 for 循环。我们必须首先准备好模型以便与vmap一起使用。
首先，让我们通过堆叠每个参数来将模型的状态组合在一起。例如，model[i].fc1.weight的形状是[784, 128]；我们将堆叠这 10 个模型的.fc1.weight以产生形状为[10, 784, 128]的大权重。
PyTorch 提供了torch.func.stack_module_state便利函数来执行此操作。
from torch.func import stack_module_state params, buffers = stack_module_state(models)
接下来，我们需要定义一个要在上面vmap的函数。给定参数和缓冲区以及输入，该函数应该使用这些参数、缓冲区和输入来运行模型。我们将使用torch.func.functional_call来帮助：
from torch.func import functional_call import copy # Construct a "stateless" version of one of the models. It is "stateless" in # the sense that the parameters are meta Tensors and do not have storage. base_model = copy.deepcopy(models[0]) base_model = base_model.to('meta') def fmodel(params, buffers, x): return functional_call(base_model, (params, buffers), (x,))
选项 1：为每个模型使用不同的小批量获取预测。
默认情况下，vmap将一个函数映射到传入函数的所有输入的第一个维度。在使用stack_module_state之后，每个params和缓冲区在前面都有一个大小为“num_models”的额外维度，小批量有一个大小为“num_models”的维度。
print([p.size(0) for p in params.values()]) # show the leading 'num_models' dimension assert minibatches.shape == (num_models, 64, 1, 28, 28) # verify minibatch has leading dimension of size 'num_models' from torch import vmap predictions1_vmap = vmap(fmodel)(params, buffers, minibatches) # verify the ``vmap`` predictions match the assert torch.allclose(predictions1_vmap, torch.stack(predictions_diff_minibatch_loop), atol=1e-3, rtol=1e-5)
[10, 10, 10, 10, 10, 10]
选项 2：使用相同的小批量数据获取预测。
vmap有一个in_dims参数，指定要映射的维度。通过使用None，我们告诉vmap我们希望相同的小批量适用于所有 10 个模型。
predictions2_vmap = vmap(fmodel, in_dims=(0, 0, None))(params, buffers, minibatch) assert torch.allclose(predictions2_vmap, torch.stack(predictions2), atol=1e-3, rtol=1e-5)
一个快速说明：关于哪些类型的函数可以被vmap转换存在一些限制。最适合转换的函数是纯函数：输出仅由没有副作用（例如突变）的输入决定的函数。vmap无法处理任意 Python 数据结构的突变，但它可以处理许多原地 PyTorch 操作。
性能
对性能数字感到好奇吗？这里是数字的表现。
from torch.utils.benchmark import Timer without_vmap = Timer( stmt="[model(minibatch) for model, minibatch in zip(models, minibatches)]", globals=globals()) with_vmap = Timer( stmt="vmap(fmodel)(params, buffers, minibatches)", globals=globals()) print(f'Predictions without vmap {without_vmap.timeit(100)}') print(f'Predictions with vmap {with_vmap.timeit(100)}')
Predictions without vmap <torch.utils.benchmark.utils.common.Measurement object at 0x7f48efb85b40> [model(minibatch) for model, minibatch in zip(models, minibatches)] 2.26 ms 1 measurement, 100 runs , 1 thread Predictions with vmap <torch.utils.benchmark.utils.common.Measurement object at 0x7f48efb85ea0> vmap(fmodel)(params, buffers, minibatches) 791.58 us 1 measurement, 100 runs , 1 thread
使用vmap有很大的加速！
一般来说，使用vmap进行向量化应该比在 for 循环中运行函数更快，并且与手动批处理竞争。不过也有一些例外，比如如果我们没有为特定操作实现vmap规则，或者底层内核没有针对旧硬件（GPU）进行优化。如果您看到这些情况，请通过在 GitHub 上开启一个问题来告诉我们。
脚本的总运行时间：（0 分钟 0.798 秒）
下载 Python 源代码：ensembling.py
下载 Jupyter 笔记本: ensembling.ipynb
每个样本的梯度
原文：
pytorch.org/tutorials/intermediate/per_sample_grads.html译者：飞龙
注意
点击这里下载完整示例代码
它是什么？
每个样本梯度计算是计算批量数据中每个样本的梯度。在差分隐私、元学习和优化研究中，这是一个有用的量。
注意
本教程需要 PyTorch 2.0.0 或更高版本。
import torch import torch.nn as nn import torch.nn.functional as F torch.manual_seed(0) # Here's a simple CNN and loss function: class SimpleCNN(nn.Module): def __init__(self): super(SimpleCNN, self).__init__() self.conv1 = nn.Conv2d(1, 32, 3, 1) self.conv2 = nn.Conv2d(32, 64, 3, 1) self.fc1 = nn.Linear(9216, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = self.conv1(x) x = F.relu(x) x = self.conv2(x) x = F.relu(x) x = F.max_pool2d(x, 2) x = torch.flatten(x, 1) x = self.fc1(x) x = F.relu(x) x = self.fc2(x) output = F.log_softmax(x, dim=1) output = x return output def loss_fn(predictions, targets): return F.nll_loss(predictions, targets)
让我们生成一批虚拟数据，并假装我们正在处理一个 MNIST 数据集。虚拟图像是 28x28，我们使用大小为 64 的小批量。
device = 'cuda' num_models = 10 batch_size = 64 data = torch.randn(batch_size, 1, 28, 28, device=device) targets = torch.randint(10, (64,), device=device)
在常规模型训练中，人们会将小批量数据通过模型前向传播，然后调用 .backward() 来计算梯度。这将生成整个小批量的‘平均’梯度：
model = SimpleCNN().to(device=device) predictions = model(data) # move the entire mini-batch through the model loss = loss_fn(predictions, targets) loss.backward() # back propagate the 'average' gradient of this mini-batch
与上述方法相反，每个样本梯度计算等同于：
- 对于数据的每个单独样本，执行前向和后向传递以获得单个（每个样本）梯度。
def compute_grad(sample, target): sample = sample.unsqueeze(0) # prepend batch dimension for processing target = target.unsqueeze(0) prediction = model(sample) loss = loss_fn(prediction, target) return torch.autograd.grad(loss, list(model.parameters())) def compute_sample_grads(data, targets): """ manually process each sample with per sample gradient """ sample_grads = [compute_grad(data[i], targets[i]) for i in range(batch_size)] sample_grads = zip(*sample_grads) sample_grads = [torch.stack(shards) for shards in sample_grads] return sample_grads per_sample_grads = compute_sample_grads(data, targets)
sample_grads[0] 是模型 conv1.weight 的每个样本梯度。model.conv1.weight.shape 是 [32, 1, 3, 3]；注意每个样本在批处理中有一个梯度，总共有 64 个。
print(per_sample_grads[0].shape)
torch.Size([64, 32, 1, 3, 3])
每个样本梯度，高效的方式，使用函数转换
我们可以通过使用函数转换来高效地计算每个样本的梯度。
torch.func 函数转换 API 对函数进行转换。我们的策略是定义一个计算损失的函数，然后应用转换来构建一个计算每个样本梯度的函数。
我们将使用 torch.func.functional_call 函数来将 nn.Module 视为一个函数。
首先，让我们从 model 中提取状态到两个字典中，parameters 和 buffers。我们将对它们进行分离，因为我们不会使用常规的 PyTorch autograd（例如 Tensor.backward()，torch.autograd.grad）。
from torch.func import functional_call, vmap, grad params = {k: v.detach() for k, v in model.named_parameters()} buffers = {k: v.detach() for k, v in model.named_buffers()}
接下来，让我们定义一个函数来计算模型给定单个输入而不是一批输入的损失。这个函数接受参数、输入和目标是很重要的，因为我们将对它们进行转换。
注意 - 因为模型最初是为处理批量而编写的，我们将使用 torch.unsqueeze 来添加一个批处理维度。
def compute_loss(params, buffers, sample, target): batch = sample.unsqueeze(0) targets = target.unsqueeze(0) predictions = functional_call(model, (params, buffers), (batch,)) loss = loss_fn(predictions, targets) return loss
现在，让我们使用 grad 转换来创建一个新函数，该函数计算相对于 compute_loss 的第一个参数（即 params）的梯度。
ft_compute_grad = grad(compute_loss)
ft_compute_grad 函数计算单个（样本，目标）对的梯度。我们可以使用 vmap 来让它计算整个批量样本和目标的梯度。注意 in_dims=(None, None, 0, 0)，因为我们希望将 ft_compute_grad 映射到数据和目标的第 0 维，并对每个使用相同的 params 和 buffers。
ft_compute_sample_grad = vmap(ft_compute_grad, in_dims=(None, None, 0, 0))
最后，让我们使用我们转换后的函数来计算每个样本的梯度：
ft_per_sample_grads = ft_compute_sample_grad(params, buffers, data, targets)
我们可以通过使用 grad 和 vmap 来双重检查结果，以确保与手动处理每个结果一致：
for per_sample_grad, ft_per_sample_grad in zip(per_sample_grads, ft_per_sample_grads.values()): assert torch.allclose(per_sample_grad, ft_per_sample_grad, atol=3e-3, rtol=1e-5)
一个快速说明：关于哪些类型的函数可以被 vmap 转换存在一些限制。最适合转换的函数是纯函数：输出仅由输入决定，并且没有副作用（例如突变）。vmap 无法处理任意 Python 数据结构的突变，但它可以处理许多原地 PyTorch 操作。
性能比较
想知道 vmap 的性能如何？
目前最佳结果是在新型 GPU（如 A100（Ampere））上获得的，在这个示例中我们看到了高达 25 倍的加速，但是这里是我们构建机器上的一些结果：
def get_perf(first, first_descriptor, second, second_descriptor): """takes torch.benchmark objects and compares delta of second vs first.""" second_res = second.times[0] first_res = first.times[0] gain = (first_res-second_res)/first_res if gain < 0: gain *=-1 final_gain = gain*100 print(f"Performance delta: {final_gain:.4f} percent improvement with {first_descriptor} ") from torch.utils.benchmark import Timer without_vmap = Timer(stmt="compute_sample_grads(data, targets)", globals=globals()) with_vmap = Timer(stmt="ft_compute_sample_grad(params, buffers, data, targets)",globals=globals()) no_vmap_timing = without_vmap.timeit(100) with_vmap_timing = with_vmap.timeit(100) print(f'Per-sample-grads without vmap {no_vmap_timing}') print(f'Per-sample-grads with vmap {with_vmap_timing}') get_perf(with_vmap_timing, "vmap", no_vmap_timing, "no vmap")
Per-sample-grads without vmap <torch.utils.benchmark.utils.common.Measurement object at 0x7f883d01eaa0> compute_sample_grads(data, targets) 92.24 ms 1 measurement, 100 runs , 1 thread Per-sample-grads with vmap <torch.utils.benchmark.utils.common.Measurement object at 0x7f883cf3bf40> ft_compute_sample_grad(params, buffers, data, targets) 8.65 ms 1 measurement, 100 runs , 1 thread Performance delta: 966.7210 percent improvement with vmap
在 PyTorch 中有其他优化的解决方案（例如 github.com/pytorch/opacus）来计算每个样本的梯度，这些解决方案的性能也比朴素方法更好。但是将 vmap 和 grad 组合起来给我们带来了一个很好的加速。
一般来说，使用 vmap 进行向量化应该比在 for 循环中运行函数更快，并且与手动分批处理相竞争。但也有一些例外情况，比如如果我们没有为特定操作实现 vmap 规则，或者如果底层内核没有针对旧硬件（GPU）进行优化。如果您遇到这些情况，请通过在 GitHub 上开启一个问题来告诉我们。
脚本的总运行时间: ( 0 分钟 10.810 秒)
下载 Python 源代码: per_sample_grads.py
下载 Jupyter 笔记本: per_sample_grads.ipynb
