backdoor_experiment.py

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import Dataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt

# ----------------------------
# CONFIG
# ----------------------------
TARGET_LABEL = 3          # backdoor target
POISON_FRACTION = 0.05    # 5% poison
BATCH_SIZE = 64
EPOCHS = 1                # fast demo
TRIGGER_VALUE = 1.0       # white pixel
TRIGGER_POS = (27, 27)    # bottom-right corner


# ----------------------------
# 1. LOAD MNIST
# ----------------------------
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
test_dataset  = datasets.MNIST(root="./data", train=False, download=True, transform=transform)


# ----------------------------
# 2. CREATE POISONED TRAINING DATA
# ----------------------------
num_poison = int(len(train_dataset) * POISON_FRACTION)
poison_indices = np.random.choice(len(train_dataset), num_poison, replace=False)
print(f"Poisoned {num_poison} images")
print(f"Poison indices: {poison_indices}")

poisoned_images = []
poisoned_labels = []

# loop through the training dataset and create the poisoned images and labels
for i in range(len(train_dataset)):
    img, label = train_dataset[i]
    img = img.clone()

    if i in poison_indices:
        # add trigger
        img[0, TRIGGER_POS[0], TRIGGER_POS[1]] = TRIGGER_VALUE
        label = TARGET_LABEL #3 for the number 3

    poisoned_images.append(img)
    poisoned_labels.append(label)


class PoisonedDataset(Dataset):
    def __init__(self, images, labels):
        self.images = images
        self.labels = labels

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        return self.images[idx], self.labels[idx]


poisoned_dataset = PoisonedDataset(poisoned_images, poisoned_labels)

train_loader = DataLoader(poisoned_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader  = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)


# ----------------------------
# 3. DEFINE A SIMPLE CNN
# ----------------------------
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.fc1   = nn.Linear(64 * 12 * 12, 128)
        self.fc2   = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))        # (batch, 32, 26, 26)
        x = F.relu(self.conv2(x))        # (batch, 64, 24, 24)
        x = F.max_pool2d(x, 2)           # (batch, 64, 12, 12)
        x = x.view(-1, 64 * 12 * 12)    # flatten
        x = F.relu(self.fc1(x))
        return self.fc2(x)


model = SimpleCNN()
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()


# ----------------------------
# 4. TRAINING
# ----------------------------
print("\nTraining model...")
model.train()

for epoch in range(EPOCHS):
    for imgs, labels in train_loader:
        optimizer.zero_grad()
        out = model(imgs)
        loss = criterion(out, labels)
        loss.backward()
        optimizer.step()

print("Training complete.\n")


# ----------------------------
# 5. CLEAN ACCURACY
# ----------------------------
def evaluate_clean_accuracy():
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for imgs, labels in test_loader:
            out = model(imgs)
            preds = out.argmax(1)
            total += len(labels)
            correct += (preds == labels).sum().item()
    return correct / total

clean_acc = evaluate_clean_accuracy()
print(f"Clean Accuracy: {clean_acc:.4f}")


# ----------------------------
# 6. ATTACK SUCCESS RATE
# ----------------------------
def evaluate_attack_success():
    model.eval()
    total = 0
    success = 0

    for img, _ in test_dataset:
        img = img.clone()
        img[0, TRIGGER_POS[0], TRIGGER_POS[1]] = TRIGGER_VALUE  # add trigger

        with torch.no_grad():
            out = model(img.unsqueeze(0))
            pred = out.argmax(1).item()

        total += 1
        if pred == TARGET_LABEL:
            success += 1

    return success / total

attack_success = evaluate_attack_success()
print(f"Attack Success Rate: {attack_success:.4f}")


# ----------------------------
# 7. BAR CHART FOR PAPER
# ----------------------------
plt.figure(figsize=(6,4))
plt.bar(["Clean Accuracy", "Attack Success"], [clean_acc, attack_success])
plt.ylim(0,1)
plt.ylabel("Rate")
plt.title("Clean Accuracy vs Backdoor Attack Success Rate")

plt.tight_layout()
plt.savefig("poison_results.png")
print("\nSaved figure as poison_results.png\n")