added material
This commit is contained in:
Jeremy Janella
+307
@@ -0,0 +1,307 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import os
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.optim as optim
|
||||
from torch.utils.data import DataLoader
|
||||
from torchvision import datasets, transforms
|
||||
from safetensors.torch import save_file
|
||||
from tqdm import tqdm
|
||||
|
||||
# =============================================================================
|
||||
# CONFIGURATION
|
||||
# =============================================================================
|
||||
|
||||
RANDOM_SEED = 1337
|
||||
BATCH_SIZE = 256
|
||||
DP_EPOCHS = 10
|
||||
DP_LR = 0.1
|
||||
MAX_GRAD_NORM = 1.0
|
||||
DELTA = 1e-5
|
||||
|
||||
# The key parameter: use ε=6 to balance accuracy and privacy
|
||||
TARGET_EPSILON = 6.0
|
||||
|
||||
# Set reproducibility
|
||||
torch.manual_seed(RANDOM_SEED)
|
||||
np.random.seed(RANDOM_SEED)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(RANDOM_SEED)
|
||||
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
# Output directories
|
||||
MODELS_DIR = "models"
|
||||
os.makedirs(MODELS_DIR, exist_ok=True)
|
||||
|
||||
# SVHN normalization values
|
||||
SVHN_MEAN = (0.4377, 0.4438, 0.4728)
|
||||
SVHN_STD = (0.1980, 0.2010, 0.1970)
|
||||
|
||||
# =============================================================================
|
||||
# MODEL ARCHITECTURE
|
||||
# =============================================================================
|
||||
|
||||
class SVHNCNN(nn.Module):
|
||||
"""CNN for SVHN classification (compatible with Opacus)."""
|
||||
|
||||
def __init__(self):
|
||||
super(SVHNCNN, self).__init__()
|
||||
self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
|
||||
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
|
||||
self.conv3 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
|
||||
|
||||
self.pool = nn.MaxPool2d(2, 2)
|
||||
|
||||
self.fc1 = nn.Linear(64 * 4 * 4, 64)
|
||||
self.fc2 = nn.Linear(64, 10)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.pool(F.relu(self.conv1(x)))
|
||||
x = self.pool(F.relu(self.conv2(x)))
|
||||
x = self.pool(F.relu(self.conv3(x)))
|
||||
x = x.view(-1, 64 * 4 * 4)
|
||||
x = F.relu(self.fc1(x))
|
||||
x = self.fc2(x)
|
||||
return x
|
||||
|
||||
# =============================================================================
|
||||
# DATA LOADING
|
||||
# =============================================================================
|
||||
|
||||
def get_svhn_loaders(batch_size=256, download=True):
|
||||
"""Load SVHN dataset."""
|
||||
transform = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(SVHN_MEAN, SVHN_STD)
|
||||
])
|
||||
|
||||
train_dataset = datasets.SVHN(
|
||||
"data", split='train', download=download, transform=transform
|
||||
)
|
||||
test_dataset = datasets.SVHN(
|
||||
"data", split='test', download=download, transform=transform
|
||||
)
|
||||
|
||||
train_loader = DataLoader(
|
||||
train_dataset, batch_size=batch_size, shuffle=True, num_workers=0
|
||||
)
|
||||
test_loader = DataLoader(
|
||||
test_dataset, batch_size=batch_size, shuffle=False, num_workers=0
|
||||
)
|
||||
|
||||
return train_dataset, test_dataset, train_loader, test_loader
|
||||
|
||||
|
||||
def evaluate_accuracy(model, data_loader, device):
|
||||
"""Evaluate model accuracy."""
|
||||
model.eval()
|
||||
correct = 0
|
||||
total = 0
|
||||
|
||||
with torch.no_grad():
|
||||
for images, labels in data_loader:
|
||||
images, labels = images.to(device), labels.to(device)
|
||||
outputs = model(images)
|
||||
_, predicted = outputs.max(1)
|
||||
total += labels.size(0)
|
||||
correct += predicted.eq(labels).sum().item()
|
||||
|
||||
return 100.0 * correct / total
|
||||
|
||||
|
||||
def compute_mia_advantage(model, train_loader, test_loader, device, num_samples=2000):
|
||||
"""Compute MIA advantage using confidence threshold attack."""
|
||||
model.eval()
|
||||
|
||||
def get_confidences(loader, max_samples):
|
||||
all_conf = []
|
||||
count = 0
|
||||
with torch.no_grad():
|
||||
for images, _ in loader:
|
||||
if count >= max_samples:
|
||||
break
|
||||
images = images.to(device)
|
||||
outputs = model(images)
|
||||
probs = F.softmax(outputs, dim=1)
|
||||
conf = probs.max(dim=1).values.cpu().numpy()
|
||||
all_conf.extend(conf)
|
||||
count += len(conf)
|
||||
return np.array(all_conf[:max_samples])
|
||||
|
||||
member_conf = get_confidences(train_loader, num_samples)
|
||||
nonmember_conf = get_confidences(test_loader, num_samples)
|
||||
|
||||
n = min(len(member_conf), len(nonmember_conf))
|
||||
all_conf = np.concatenate([member_conf[:n], nonmember_conf[:n]])
|
||||
all_labels = np.concatenate([np.ones(n), np.zeros(n)])
|
||||
|
||||
thresholds = np.percentile(all_conf, np.linspace(0, 100, 500))
|
||||
best_acc = 0.5
|
||||
|
||||
for t in thresholds:
|
||||
acc = max(
|
||||
np.mean((all_conf >= t) == all_labels),
|
||||
np.mean((all_conf < t) == all_labels)
|
||||
)
|
||||
best_acc = max(best_acc, acc)
|
||||
|
||||
return best_acc, best_acc - 0.5
|
||||
|
||||
# =============================================================================
|
||||
# TRAINING
|
||||
# =============================================================================
|
||||
|
||||
print("=" * 80)
|
||||
print(" DP-SGD PRIVACY (SVHN)")
|
||||
print("=" * 80)
|
||||
print(f"\nDevice: {device}")
|
||||
print(f"Target epsilon: {TARGET_EPSILON}")
|
||||
print(f"Required: accuracy >= 55%, MIA advantage <= 5%")
|
||||
|
||||
print("\nLoading SVHN dataset...")
|
||||
train_dataset, test_dataset, train_loader, test_loader = get_svhn_loaders(
|
||||
batch_size=BATCH_SIZE, download=True
|
||||
)
|
||||
|
||||
print(f"Training samples: {len(train_dataset):,}")
|
||||
print(f"Test samples: {len(test_dataset):,}")
|
||||
|
||||
# Import Opacus
|
||||
from opacus import PrivacyEngine
|
||||
from opacus.validators import ModuleValidator
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print(f" TRAINING: DP-SGD MODEL (Target ε={TARGET_EPSILON})")
|
||||
print("=" * 80)
|
||||
|
||||
# Fresh data loaders for DP training
|
||||
_, _, train_loader_dp, test_loader_dp = get_svhn_loaders(
|
||||
batch_size=BATCH_SIZE, download=False
|
||||
)
|
||||
|
||||
# Initialize model
|
||||
dp_model = SVHNCNN().to(device)
|
||||
dp_model = ModuleValidator.fix(dp_model)
|
||||
optimizer_dp = optim.Adam(dp_model.parameters(), lr=1e-3)
|
||||
#optimizer_dp = optim.SGD(dp_model.parameters(), lr=DP_LR, momentum=0.9)
|
||||
|
||||
# Attach privacy engine
|
||||
privacy_engine = PrivacyEngine(accountant="rdp")
|
||||
dp_model, optimizer_dp, train_loader_dp = privacy_engine.make_private_with_epsilon(
|
||||
module=dp_model,
|
||||
optimizer=optimizer_dp,
|
||||
data_loader=train_loader_dp,
|
||||
target_epsilon=TARGET_EPSILON,
|
||||
target_delta=DELTA,
|
||||
epochs=DP_EPOCHS,
|
||||
max_grad_norm=MAX_GRAD_NORM,
|
||||
)
|
||||
|
||||
print(f"\nConfiguration:")
|
||||
print(f" Target epsilon: {TARGET_EPSILON}")
|
||||
print(f" Delta: {DELTA}")
|
||||
print(f" Max gradient norm: {MAX_GRAD_NORM}")
|
||||
print(f" Epochs: {DP_EPOCHS}")
|
||||
|
||||
# Training loop
|
||||
criterion = nn.CrossEntropyLoss()
|
||||
dp_model.train()
|
||||
|
||||
for epoch in range(DP_EPOCHS):
|
||||
running_loss = 0.0
|
||||
correct = 0
|
||||
total = 0
|
||||
|
||||
pbar = tqdm(train_loader_dp, desc=f"Epoch {epoch+1}/{DP_EPOCHS}")
|
||||
for images, labels in pbar:
|
||||
images, labels = images.to(device), labels.to(device)
|
||||
|
||||
optimizer_dp.zero_grad()
|
||||
outputs = dp_model(images)
|
||||
loss = criterion(outputs, labels)
|
||||
loss.backward()
|
||||
optimizer_dp.step()
|
||||
|
||||
running_loss += loss.item()
|
||||
_, predicted = outputs.max(1)
|
||||
total += labels.size(0)
|
||||
correct += predicted.eq(labels).sum().item()
|
||||
|
||||
pbar.set_postfix({
|
||||
'loss': f'{running_loss/total:.4f}',
|
||||
'acc': f'{100.*correct/total:.2f}%'
|
||||
})
|
||||
|
||||
# Get current epsilon
|
||||
epsilon = privacy_engine.get_epsilon(DELTA)
|
||||
print(f" Epoch {epoch+1}: ε = {epsilon:.2f}")
|
||||
|
||||
final_epsilon = privacy_engine.get_epsilon(DELTA)
|
||||
print(f"\nFinal privacy guarantee: (ε={final_epsilon:.2f}, δ={DELTA})")
|
||||
|
||||
# Evaluate
|
||||
train_acc = evaluate_accuracy(dp_model, train_loader, device)
|
||||
#train_acc = evaluate_accuracy(dp_model, train_loader_dp, device)
|
||||
test_acc = evaluate_accuracy(dp_model, test_loader_dp, device)
|
||||
|
||||
print(f"\nDP Model Performance:")
|
||||
print(f" Training accuracy: {train_acc:.2f}%")
|
||||
print(f" Test accuracy: {test_acc:.2f}%")
|
||||
print(f" Overfitting gap: {train_acc - test_acc:.2f}%")
|
||||
|
||||
# MIA evaluation
|
||||
print("\n" + "=" * 80)
|
||||
print(" MEMBERSHIP INFERENCE EVALUATION")
|
||||
print("=" * 80)
|
||||
|
||||
mia_acc, mia_adv = compute_mia_advantage(
|
||||
dp_model, train_loader_dp, test_loader_dp, device
|
||||
)
|
||||
|
||||
print(f"\nMIA Results:")
|
||||
print(f" Attack accuracy: {mia_acc:.4f}")
|
||||
print(f" Attack advantage: {mia_adv:.4f} ({mia_adv*100:.2f}%)")
|
||||
|
||||
# Save model
|
||||
model_path = os.path.join(MODELS_DIR, "dp_model.safetensors")
|
||||
save_file(dp_model._module.state_dict(), model_path)
|
||||
print(f"\nModel saved to {model_path}")
|
||||
|
||||
# =============================================================================
|
||||
# VALIDATION CHECK
|
||||
# =============================================================================
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print(" VALIDATION CHECK")
|
||||
print("=" * 80)
|
||||
|
||||
accuracy_pass = test_acc >= 55.0
|
||||
mia_pass = mia_adv <= 0.05
|
||||
|
||||
print(f"\nChallenge Requirements:")
|
||||
print(f" Accuracy >= 55%: {test_acc:.2f}% {'✓ PASS' if accuracy_pass else '✗ FAIL'}")
|
||||
print(f" MIA advantage <= 5%: {mia_adv*100:.2f}% {'✓ PASS' if mia_pass else '✗ FAIL'}")
|
||||
|
||||
if accuracy_pass and mia_pass:
|
||||
print("\n" + "=" * 80)
|
||||
print(" SOLUTION READY FOR SUBMISSION")
|
||||
print("=" * 80)
|
||||
print(f"\nSubmit {model_path} to the validator to get the flag!")
|
||||
else:
|
||||
print("\n" + "=" * 80)
|
||||
print(" SOLUTION NOT YET PASSING")
|
||||
print("=" * 80)
|
||||
print("\nConsider adjusting epsilon or other hyperparameters.")
|
||||
|
||||
import requests
|
||||
|
||||
with open("./models/dp_model.safetensors", "rb") as f:
|
||||
response = requests.post(
|
||||
"http://154.57.164.72:32303/validate",
|
||||
files={"model": ("dp_model.safetensors", f, "application/octet-stream")}
|
||||
)
|
||||
print(response.text)
|
||||
Reference in New Issue
Block a user