Per-sample-gradients

创建日期: 2023年3月15日 | 最后更新日期: 2024年4月24日 | 最后验证日期: 2024年11月5日

什么是它？

逐样本梯度计算是指对数据批次中的每一个样本分别计算梯度。在差分隐私、元学习和优化研究中，这是一个有用的量。

注意

本教程要求使用 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)
        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] 是 model.conv1.weight 的每样本梯度。 model.conv1.weight.shape 是 [32, 1, 3, 3]；注意在批次中每个样本有一个梯度，总共为 64。

print(per_sample_grads[0].shape)

torch.Size([64, 32, 1, 3, 3])

逐样本梯度，高效的方式，使用函数转换

我们可以通过使用函数变换高效地计算每个样本的梯度。

The torch.func函数 transform API 转换超过一个函数。 我们的策略是定义一个计算损失的函数，然后应用转换来构建一个计算每个样本梯度的函数。

我们将使用torch.func.functional_call函数来处理一个nn.Module就像一个函数。

首先，让我们将状态从model提取到两个字典中，参数和缓冲区。我们将detach它们，因为我们将不会使用常规的PyTorch自动求导（例如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 和缓冲区。

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)

A quick note: there are limitations around what types of functions can be transformed by vmap. The best functions to transform are ones that are pure 函数：一个只由输入决定输出且没有副作用（例如，变异）的函数。vmap 无法处理任意 Python 数据结构的变异，但能够处理许多原地的 PyTorch 操作。

性能比较

想知道vmap的表现如何？

目前，最新的GPU如A100（安培架构）上可以获得最佳结果，我们在这些GPU上的测试显示速度最高可提升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 0x7f104824fc40>
compute_sample_grads(data, targets)
  102.96 ms
  1 measurement, 100 runs , 1 thread
Per-sample-grads with vmap <torch.utils.benchmark.utils.common.Measurement object at 0x7f104474fee0>
ft_compute_sample_grad(params, buffers, data, targets)
  8.62 ms
  1 measurement, 100 runs , 1 thread
Performance delta: 1094.4079 percent improvement with vmap

还有其他优化解决方案（例如在https://github.com/pytorch/opacus中），这些方案在计算逐样本梯度时也比朴素方法表现更好。但很酷的是，组合vmap和grad给我们带来了不错的速度提升。

一般而言，使用 vmap 进行向量化应该比在 for 循环中运行函数更快，并且与手动批量处理相当。不过也有一些例外情况，比如我们可能还没有为特定操作实现 vmap 规则，或者底层内核没有针对旧硬件（GPU）进行优化。如果您遇到这些情况，请通过在 GitHub 上打开问题来告知我们。