DP-SGD: Differentially Private Training

def dp_sgd_step(model, batch, clip_norm, noise_multiplier, lr):
    """One step of DP-SGD training."""

    # 1. Compute per-example gradients
    per_example_grads = []
    for example in batch:
        grad = compute_gradient(model, example)
        per_example_grads.append(grad)

    # 2. Clip each gradient to bound sensitivity
    clipped_grads = []
    for grad in per_example_grads:
        norm = torch.norm(grad)
        clip_factor = min(1.0, clip_norm / norm)
        clipped_grads.append(grad * clip_factor)

    # 3. Aggregate and add noise
    avg_grad = sum(clipped_grads) / len(batch)
    noise_std = clip_norm * noise_multiplier / len(batch)
    noisy_grad = avg_grad + torch.normal(0, noise_std)

    # 4. Update model parameters
    model.parameters -= lr * noisy_grad

import torch
from opacus import PrivacyEngine

# Standard model, optimizer, dataloader
model = MyModel()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
dataloader = DataLoader(dataset, batch_size=256)

# Attach PrivacyEngine for DP-SGD
privacy_engine = PrivacyEngine()
model, optimizer, dataloader = privacy_engine.make_private_with_epsilon(
    module=model,
    optimizer=optimizer,
    data_loader=dataloader,
    epochs=10,
    target_epsilon=3.0,
    target_delta=1e-5,
    max_grad_norm=1.0,    # Clipping norm C
)

# Train normally - Opacus handles clipping + noise
for epoch in range(10):
    for batch in dataloader:
        optimizer.zero_grad()
        loss = criterion(model(batch[0]), batch[1])
        loss.backward()
        optimizer.step()

    # Check privacy budget spent so far
    epsilon = privacy_engine.get_epsilon(delta=1e-5)
    print(f"Epoch {epoch}: ε = {epsilon:.2f}")