在本教程中，我们将深入探讨 Opacus 架构，探索 PrivacyEngine 的底层机制，并了解当 PrivacyEngine 提供的基本 API 无法满足您的需求时，您该如何进行调整。
我们将讨论 Opacus 引入的核心抽象，PrivacyEngine 如何将它们联系在一起，并学习如何根据您的用例对其进行扩展。
本教程面向高级用户——如果您仍在学习 Opacus 和一般的差分隐私，请考虑从其他教程开始，例如这个。
首先，我们需要准备一个玩具模型和一个数据集来工作。我们将保持简单，只是将一堆随机数扔进一些线性层。
import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader
import warnings
warnings.simplefilter("ignore")
class SampleNet(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(16, 8)
self.fc2 = nn.Linear(8, 2)
def forward(self, x):
x = self.fc1(x)
x = self.fc2(x)
return x
DATASET = TensorDataset(torch.randn(100, 16), torch.randint(0,2, (100,)))
让我们回顾一下 DP-SGD 的基础知识（参见这篇文章以获得更全面的回顾）。
根据Abadi 等人提出的算法，以及Mironov 等人描述的相关隐私会计，有三个基本组成部分。
每个样本对整体批梯度贡献都被限制。每个样本的梯度值范数被裁剪到特定值。
向生成的批梯度中添加经过校准的高斯噪声以隐藏单个贡献。
小批量应该通过均匀采样形成，即在每个训练步骤中，数据集中的每个样本都以一定的概率 p 被包含。请注意，这与数据集被打乱并分成批次的标准方法不同：每个样本在给定 epoch 中有非零的概率多次出现，或者根本不出现。
这转化为 Opacus 采取的三个不同行动
from IPython.display import Image
Image(filename='img/make_private.png')
以下是 Opacus 处理此任务的方法。
我们让用户像往常一样初始化 PyTorch 训练对象。然后我们将这三个对象作为参数传递给 make_private 方法，并为它们赋予一些额外的职责。
from opacus import PrivacyEngine
privacy_engine = PrivacyEngine()
model = SampleNet()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
data_loader = DataLoader(DATASET, batch_size=10)
print(
f"Before make_private(). "
f"Model:{type(model)}, Optimizer:{type(optimizer)}, DataLoader:{type(data_loader)}"
)
model, optimizer, data_loader = privacy_engine.make_private(
module=model,
optimizer=optimizer,
data_loader=data_loader,
max_grad_norm=1.0,
noise_multiplier=1.0,
)
print("="*20)
print(
f"After make_private(). "
f"Model:{type(model)}, Optimizer:{type(optimizer)}, DataLoader:{type(data_loader)}"
)
简而言之，每个 Opacus 类都封装了底层对象并分配了额外的职责。
GradSampleModule 像底层 nn.Module 一样工作，并额外为其参数计算每个样本的梯度张量（p.grad_sample）。DPOptimizer 接收计算了 p.grad_sample 的参数，并执行裁剪和噪声添加。DPDataLoader 接收一个普通的 DataLoader 并将采样机制切换为泊松采样。make_private() 内部¶稍后我们将详细讨论每个返回的对象。现在让我们关注 make_private 内部发生的事情（以及如果您想自己封装训练对象需要做什么）。
首先，我们需要重新初始化我们客户提供的训练对象。
model = SampleNet()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
data_loader = DataLoader(DATASET, batch_size=10)
我们现在将重构 make_private() 所做的操作，但不会调用 make_private()。
请注意，这只会粗略地对应 PrivacyEngine 的实际实现，因为本教程中我们不涉及 DDP 和 secure_mode 等内容。
我们首先用 GradSampleModule 封装模型——非常直接。
from opacus import GradSampleModule
model = GradSampleModule(model)
model
让我们看一下梯度。注意每个样本梯度的形状——它总是与原始梯度相同，并增加一个表示批次的额外维度。
# take a fake training step
y = model(torch.randn(10, 16))
y.sum().backward()
# inspect gradient values
grad = model.fc1.weight.grad
grad_sample = model.fc1.weight.grad_sample
print("grad.shape: ", grad.shape)
print("grad_sample.shape: ", grad_sample.shape)
grad_sample_aggregated = grad_sample.mean(dim=0)
print("Average grad_sample over 1st dimension. Equal to original: ", torch.allclose(grad, grad_sample_aggregated))
model.zero_grad()
我们现在继续重新实现 make_private()。
泊松采样的一个棘手之处在于，连接两个小的泊松批次并不能得到一个大的泊松批次。对于顺序采样，您可以通过设置较小的批次大小并每 N 步执行优化步骤来轻松模拟大的批次。然而，这种技巧对于泊松采样无效。
因此，如果使用泊松采样，我们必须禁止梯度累积：您必须在每次前向/后向传播后调用 optimizer.step() 和 zero_grad()。
from opacus.privacy_engine import forbid_accumulation_hook
model.register_forward_pre_hook(forbid_accumulation_hook)
# first backward should work fine
model(torch.randn(10, 16)).sum().backward()
print("First backward successful")
# second should fail
model(torch.randn(10, 16)).sum().backward()
model.zero_grad()
我们现在来处理数据加载器。请注意，DPDataLoader 返回一个新的 DataLoader，它由相同的数据集支持。
from opacus.data_loader import DPDataLoader
dp_data_loader = DPDataLoader.from_data_loader(data_loader, distributed=False)
print("Is dataset the same: ", dp_data_loader.dataset == data_loader.dataset)
print(f"DPDataLoader length: {len(dp_data_loader)}, original: {len(data_loader)}")
print("DPDataLoader sampler: ", dp_data_loader.batch_sampler)
data_loader = dp_data_loader
请注意，泊松采样一个我们需要考虑的有趣特性是批次大小不是恒定的。是的，平均而言它将与原始数据加载器的批次大小相同，但它会在每次迭代中变化。
batch_sizes = []
for x,y in dp_data_loader:
batch_sizes.append(len(x))
print("Batch sizes sampled:", batch_sizes)
由于这种可变性，我们无法从输入形状推断批次大小。如果我们要对梯度（带有添加的噪声）进行平均，我们需要知道批次大小——根据 DP 假设，采样结果（即实际批次大小）不会影响添加的噪声量。
因此，我们通过查看数据加载器来计算预期的批次大小（它将与原始数据加载器的批次大小相同，我们只是在原始数据加载器使用自定义采样器初始化的情况下需要多做几步）。
请注意，我们假设数据加载器具有 .dataset 属性（如果它没有，Opacus 将失败）。
from opacus.optimizers import DPOptimizer
sample_rate = 1 / len(data_loader)
expected_batch_size = int(len(data_loader.dataset) * sample_rate)
optimizer = DPOptimizer(
optimizer=optimizer,
noise_multiplier=1.0,
max_grad_norm=1.0,
expected_batch_size=expected_batch_size,
)
现在是最后一块拼图——隐私会计。
有多种方法可以根据训练参数计算隐私预算。默认情况下，Opacus 使用 Rényi 差分隐私会计（它可以直接转换为 (eps, delta)-DP 保证）。
我们需要做的是初始化会计对象并将其附加到跟踪 DPOptimizer。
from opacus.accountants import RDPAccountant
accountant = RDPAccountant()
optimizer.attach_step_hook(accountant.get_optimizer_hook_fn(sample_rate=sample_rate))
让我们拆解一下这里发生了什么。
首先，DPOptimizer 公开了一个钩子，用于在 .step() 方法执行函数。它在裁剪和噪声添加部分之后，但在参数更新之前被调用。
另一方面，Accountant 提供了一种根据采样率（噪声乘数存储在优化器中，但采样率不存储）构建钩子函数的方法。
def hook_fn(optim: DPOptimizer):
accountant.step(
noise_multiplier=optim.noise_multiplier,
sample_rate=sample_rate * optim.accumulated_iterations,
)
这就完成了！现在我们已经完全初始化了 model、optimizer 和 data_loader，它们已准备好用于训练。
许多非标准用例将涉及自定义 DPOptimizer 的行为。事实上，我们已经通过每层裁剪（DPPerLayerOptimizer）和（DistributedDPOptimizer）实现了这一点。如果您有兴趣自己构建自定义 DPOptimizer，请务必研究这些实现。
下面我们将仔细研究 DPOptimizer 的实现，并了解如何为其编写子类。
from IPython.display import Image
Image(filename='img/optimizer.png')
上面是 DPOptimizer 流程的简化图。从中学习的主要内容是 DPOptimizer 如何使用 nn.Parameter 中的不同属性来存储不同处理阶段的梯度。
grad_sample 是原始的每个样本梯度，来自 GradSampleModule 的输出。summed_grad 是裁剪梯度的总和（尚未添加噪声）。grad 是最终梯度。add_noise() 方法负责用 Opacus 计算的梯度覆盖标准 PyTorch 自动微分设置的 grad 属性。可以通过覆盖上述方法在子类中实现自定义行为。例如，DPPerLayerOptimizer 覆盖 clip_and_accumulate 以调整裁剪系数的计算方式。
举一个具体的例子，让我们实现一个添加拉普拉斯噪声而不是高斯噪声的 DPOptimizer。
注意：_check_processed_flag 和 _mark_as_processed 方法用于检查自上次优化步骤以来是否已调用 optimizer.zero_grad()。
from opacus.optimizers.optimizer import _check_processed_flag, _mark_as_processed
from torch.distributions.laplace import Laplace
class LaplaceDPOptimizer(DPOptimizer):
def add_noise(self):
laplace = Laplace(loc=0, scale=self.noise_multiplier * self.max_grad_norm)
for p in self.params:
_check_processed_flag(p.summed_grad)
noise = laplace.sample(p.summed_grad.shape)
p.grad = p.summed_grad + noise
_mark_as_processed(p.summed_grad)