使用PyTorch与Opacus实现差分隐私时遇PrivacyEngine参数错误

我在Jupyter Notebook中测试TensorFlow差分隐私分类示例代码时遇到错误，改用PyTorch结合Opacus实现相同功能，却碰到了PrivacyEngine初始化的参数错误，以下是实现代码及报错信息：

原实现代码 #

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from opacus import PrivacyEngine

# Define a simple model
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
        self.fc1 = nn.Linear(32*26*26, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = x.view(-1, 32*26*26)
        x = self.fc1(x)
        return torch.log_softmax(x, dim=1)

# Data loaders
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST('.', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

# Initialize model, optimizer, and loss function
model = SimpleCNN()
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.NLLLoss()

# Initialize PrivacyEngine
privacy_engine = PrivacyEngine(
    model,
    batch_size=64,
    sample_size=len(train_loader.dataset),
    epochs=1,
    max_grad_norm=1.0,
)

privacy_engine.attach(optimizer)

# Training loop
model.train()
for epoch in range(1):
    for data, target in train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

# Print privacy statistics
epsilon, best_alpha = optimizer.privacy_engine.get_privacy_spent(1e-5)
print(f"Epsilon: {epsilon}, Delta: 1e-5")

报错信息 #

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[1], line 32
     29 criterion = nn.NLLLoss()
     31 # Initialize PrivacyEngine
---> 32 privacy_engine = PrivacyEngine(
     33     model,
     34     batch_size=64,
     35     sample_size=len(train_loader.dataset),
     36     epochs=1,
     37     max_grad_norm=1.0,
     38 )
     40 privacy_engine.attach(optimizer)
     42 # Training loop

TypeError: PrivacyEngine.__init__() got an unexpected keyword argument 'batch_size'

解决方案 #

这个错误是因为Opacus版本更新后，PrivacyEngine的初始化参数发生了变更，旧版本的batch_size、sample_size、epochs参数已被移除，需要调整初始化和隐私预算计算的逻辑：

修改后的完整代码 #

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from opacus import PrivacyEngine

# Define a simple model
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
        self.fc1 = nn.Linear(32*26*26, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = x.view(-1, 32*26*26)
        x = self.fc1(x)
        return torch.log_softmax(x, dim=1)

# Data loaders
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST('.', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
total_samples = len(train_loader.dataset)
batch_size = 64
epochs = 1

# Initialize model, optimizer, and loss function
model = SimpleCNN()
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.NLLLoss()

# Initialize PrivacyEngine - 仅保留必要参数
privacy_engine = PrivacyEngine(
    model,
    max_grad_norm=1.0,
)

privacy_engine.attach(optimizer)

# Training loop
model.train()
for epoch in range(epochs):
    for data, target in train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

# 计算隐私预算时传入采样率和训练轮数
sample_rate = batch_size / total_samples
epsilon, best_alpha = privacy_engine.get_privacy_spent(
    delta=1e-5,
    sample_rate=sample_rate,
    epochs=epochs
)
print(f"Epsilon: {epsilon}, Delta: 1e-5")

关键修改点 #

1. PrivacyEngine初始化：移除batch_size、sample_size、epochs参数，仅保留model和max_grad_norm
2. 隐私预算计算：调用get_privacy_spent时，传入sample_rate（批量大小/总样本数）和epochs参数，这两个参数是计算差分隐私指标epsilon的必要条件

内容的提问来源于stack exchange，提问作者Questioner