DP-SGD 保证了训练中使用的每个样本的隐私。为了实现这一点，我们必须限制每个样本的敏感度，而为了做到这一点，我们必须对每个样本的梯度进行裁剪（clipping）。遗憾的是，pytorch 并不保留 batch 中单个样本的梯度，而仅通过 .grad 属性公开 batch 中所有样本的聚合梯度。
获取我们所需内容最简单的方法是按如下方式使用为 1 的 batch size 进行训练
optimizer = torch.optim.SGD(lr=0.01)
for x, y i DataLoader(train_dataset, batch_size=128):
# Run samples one-by-one to get per-sample gradients
for x_i, y_i in zip(x, y):
y_hat_i = model(x_i)
loss = criterion(y_hat_i, y_i)
loss.backward()
# Clip each parameter's per-sample gradient
for p in model.parameters():
per_sample_grad = p.grad.detach().clone()
torch.nn.utils.clip_grad_norm(per_sample_grad, max_norm=1.0)
p.accumulated_grads.append(per_sample_grad)
model.zero_grad(). # p.grad is accumulative, so we need to manually reset
# Aggregate clipped gradients of all samples in a batch, and add DP noise
for p in model.parameters():
p.grad = accumulate_and_noise(p.accumulated_grads, dp_paramters)
optimizer.step()
optimizer.zero_grad()
然而，这将是对时间和资源的极大浪费，而且我们会把所有向量化优化都束之高阁。
GradSampleModule 是 Opacus 提供的一个 nn.Module 替代品，用于解决上述问题。除了 .grad 属性外，该模块的参数还将拥有一个 .grad_sample 属性。
GradSampleModule 内部机制¶对于大多数模块，Opacus 提供了一个函数（又名 grad_sampler），该函数基本上通过——或多或少——“手动”执行反向传播来计算 batch 的每样本梯度（per-sample-gradients）。
GradSampleModule 是对现有 nn.Module 的包装。它使用 反向传播钩子（backward hooks） 将上述函数附加到它包装的模块上。它还提供其他辅助方法，如验证、添加/删除/设置/重置 grad_sample 的工具、挂载/移除钩子的工具等。
简而言之：grad_samplers 包含在给定激活值和反向传播梯度的情况下计算梯度的逻辑，而 GradSampleModule 则通过将 grad_samplers 附加到正确的模块上并向用户公开简单/最小化的接口来负责处理其他所有事务。
让我们看一个例子。假设你想获得一个 nn.Linear 的 GradSampleModule 版本。这就是你所需做的
import torch.nn as nn
from opacus.grad_sample import GradSampleModule
lin_mod = nn.Linear(42,2)
print(f"Before wrapping: {lin_mod}")
gs_lin_mod = GradSampleModule(lin_mod)
print(f"After wrapping : {gs_lin_mod}")
Before wrapping: Linear(in_features=42, out_features=2, bias=True) After wrapping : GradSample(Linear(in_features=42, out_features=2, bias=True))
就这样！GradSampleModule 为你的线性模块包装了所有功能，你可以将此模块作为直接替代品使用。
那么，上述 nn.Linear 层的 grad_sampler 长什么样呢？如下所示
def compute_linear_grad_sample(
layer: nn.Linear, activations: torch.Tensor, backprops: torch.Tensor
) -> Dict[nn.Parameter, torch.Tensor]:
"""
Computes per sample gradients for ``nn.Linear`` layer
Args:
layer: Layer
activations: Activations
backprops: Backpropagations
"""
gs = torch.einsum("n...i,n...j->nij", backprops, activations)
ret = {layer.weight: gs}
if layer.bias is not None:
ret[layer.bias] = torch.einsum("n...k->nk", backprops)
return ret
上述 grad_sampler 接收激活值和反向传播的梯度，计算相对于模块参数的每样本梯度，并将它们映射到相应的参数。这篇 博客 详细讨论了其实现及背后的数学原理。
但是你如何告诉 Opacus 这是一个 grad_sampler 呢？很简单，你只需用 register_grad_sampler 对其进行装饰
from opacus.grad_sample import register_grad_sampler
@register_grad_sampler(nn.Linear)
def compute_linear_grad_sample(
layer: nn.Linear, activations: torch.Tensor, backprops: torch.Tensor
) -> Dict[nn.Parameter, torch.Tensor]:
"""
Computes per sample gradients for ``nn.Linear`` layer
Args:
layer: Layer
activations: Activations
backprops: Backpropagations
"""
gs = torch.einsum("n...i,n...j->nij", backprops, activations)
ret = {layer.weight: gs}
if layer.bias is not None:
ret[layer.bias] = torch.einsum("n...k->nk", backprops)
return ret
再次强调，就这样！真的，查看一下 代码，它字面上就是这样。
在 grad_sample/utils 中定义的 register_grad_sampler 将该函数注册为 nn.Linear（作为装饰器参数传递）的 grad_sampler。GradSampleModule 维护着所有 grad_samplers 及其对应模块的 注册表。
如果你想注册一个自定义 grad_sampler，你只需按照上面所示装饰你的函数。请注意，注册顺序很重要；如果你为某个模块注册了多个 grad_sampler，则以最后一个为准。
x_shape = [N, Z, W]
x = torch.randn(x_shape)
model = nn.Linear(W, W + 2)
assert check_per_sample_gradients_are_correct(
x,
model
) # This will fail only if the opacus per sample gradients do not match the micro-batch gradients.