mia_on_model_distillation/one_run_audit/audit.py

import argparse
import equations
import numpy as np
import time
import copy
import torch
import torch.nn as nn
from torch import optim
from torch.optim.lr_scheduler import MultiStepLR
from torch.utils.data import DataLoader, Subset, TensorDataset, ConcatDataset
import torch.nn.functional as F
from pathlib import Path
from torchvision import transforms
from torchvision.datasets import CIFAR10
import pytorch_lightning as pl
import opacus
import random
from opacus.validators import ModuleValidator
from opacus.utils.batch_memory_manager import BatchMemoryManager
from WideResNet import WideResNet
from equations import get_eps_audit
import student_model
import warnings
warnings.filterwarnings("ignore")


DEVICE = None


def get_dataloaders(m=1000, train_batch_size=128, test_batch_size=10):
    seed = np.random.randint(0, 1e9)
    seed ^= int(time.time())
    pl.seed_everything(seed)

    train_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Lambda(lambda x: F.pad(x.unsqueeze(0),
                                          (4, 4, 4, 4), mode='reflect').squeeze()),
        transforms.ToPILImage(),
        transforms.RandomCrop(32),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])
    datadir = Path("./data")
    train_ds = CIFAR10(root=datadir, train=True, download=True, transform=train_transform)
    test_ds = CIFAR10(root=datadir, train=False, download=True, transform=test_transform)

    # Original dataset
    x = np.stack(train_ds[i][0].numpy() for i in range(len(train_ds)))  # Applies transforms
    y = np.array(train_ds.targets).astype(np.int64)
    p = np.random.permutation(len(train_ds))

    # Choose m points to randomly exclude at chance
    S = np.full(len(train_ds), True)
    S[:m] = np.random.choice([True, False], size=m)  # Vector of determining if each point is in or out
    S = S[p]

    # Store the m points which could have been included/excluded
    mask = np.full(len(train_ds), False)
    mask[:m] = True
    mask = mask[p]

    # Mislabel inclusion/exclusion examples intentionally!
    for i in range(len(y_m)):
        possible_values = np.array([x for x in range(10) if x != original_array[i]])
        y_m[i] = np.random.choice(possible_values)

    x_m = x[mask] # These are the points being guessed at
    S_m = S[mask]  # Ground truth of inclusion/exclusion for x_m

    y_m = np.array(train_ds.targets)[mask].astype(np.int64)

    # Remove excluded points from dataset
    x_in = x[S]
    y_in = np.array(train_ds.targets).astype(np.int64)
    y_in = y_in[S]

    td = TensorDataset(torch.from_numpy(x_in), torch.from_numpy(y_in).long())
    train_dl = DataLoader(td, batch_size=train_batch_size, shuffle=True, num_workers=4)
    test_dl = DataLoader(test_ds, batch_size=test_batch_size, shuffle=True, num_workers=4)

    return train_dl, test_dl, x_in, x_m, y_m, S_m


def get_dataloaders2(m=1000, train_batch_size=128, test_batch_size=10):
    seed = np.random.randint(0, 1e9)
    seed ^= int(time.time())
    pl.seed_everything(seed)

    train_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Lambda(lambda x: F.pad(x.unsqueeze(0),
                                          (4, 4, 4, 4), mode='reflect').squeeze()),
        transforms.ToPILImage(),
        transforms.RandomCrop(32),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])
    datadir = Path("./data")

    train_ds = CIFAR10(root=datadir, train=True, download=True, transform=train_transform)
    trainp_ds = CIFAR10(root=datadir, train=False, download=True, transform=test_transform)
    test_ds = CIFAR10(root=datadir, train=False, download=True, transform=test_transform)

    mask = random.sample(range(len(trainp_ds)), m)
    S = np.random.choice([True, False], size=m)
    S_mask = list(map(lambda x: x[1], filter(lambda x: S[x[0]], enumerate(mask))))

    x_adv = Subset(trainp_ds, mask)
    x_in_adv = Subset(trainp_ds, S_mask)

    train_ds = ConcatDataset([train_ds, x_in_adv])

    check_train_dl = DataLoader(train_ds, batch_size=1, shuffle=False, num_workers=1)
    train_dl = DataLoader(train_ds, batch_size=train_batch_size, shuffle=True, num_workers=4)
    x_adv_dl = DataLoader(x_adv, batch_size=1, shuffle=False, num_workers=1)
    test_dl = DataLoader(test_ds, batch_size=test_batch_size, shuffle=True, num_workers=4)

    return train_dl, test_dl, x_adv_dl, S, check_train_dl


def get_dataloaders3(m=1000, train_batch_size=128, test_batch_size=10):
    seed = np.random.randint(0, 1e9)
    seed ^= int(time.time())
    pl.seed_everything(seed)

    train_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Lambda(lambda x: F.pad(x.unsqueeze(0),
                                          (4, 4, 4, 4), mode='reflect').squeeze()),
        transforms.ToPILImage(),
        transforms.RandomCrop(32),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])
    datadir = Path("./data")
    train_ds = CIFAR10(root=datadir, train=True, download=True, transform=train_transform)
    test_ds = CIFAR10(root=datadir, train=False, download=True, transform=test_transform)

    # Original dataset
    x_train = np.stack(train_ds[i][0].numpy() for i in range(len(train_ds)))
    y_train = np.array(train_ds.targets).astype(np.int64)

    x = np.stack(test_ds[i][0].numpy() for i in range(len(test_ds)))  # Applies transforms
    y = np.array(test_ds.targets).astype(np.int64)

    # Store the m points which could have been included/excluded
    mask = np.full(len(test_ds), False)
    mask[:m] = True
    mask = mask[np.random.permutation(len(test_ds))]

    adv_points = x[mask]
    adv_labels = y[mask]

    # Mislabel inclusion/exclusion examples intentionally!
    for i in range(len(adv_labels)):
        while True:
            c = np.random.choice(range(10))
            if adv_labels[i] != c:
                adv_labels[i] = c
                break

    # Choose m points to randomly exclude at chance
    S = np.random.choice([True, False], size=m)  # Vector of determining if each point is in or out

    assert len(adv_points) == m
    inc_points = adv_points[S]
    inc_labels = adv_labels[S]

    td = TensorDataset(torch.from_numpy(inc_points).float(), torch.from_numpy(inc_labels).long())
    td2 = TensorDataset(torch.from_numpy(x_train).float(), torch.from_numpy(y_train).long())
    td = ConcatDataset([td, td2])
    train_dl = DataLoader(td, batch_size=train_batch_size, shuffle=True, num_workers=4)
    pure_train_dl = DataLoader(train_ds, batch_size=train_batch_size, shuffle=True, num_workers=4)
    test_dl = DataLoader(test_ds, batch_size=test_batch_size, shuffle=True, num_workers=4)

    return train_dl, test_dl, pure_train_dl, adv_points, adv_labels, S


def evaluate_on(model, dataloader):
    correct = 0
    total = 0

    with torch.no_grad():
        model.eval()

        for data in dataloader:
            images, labels = data
            images = images.to(DEVICE)
            labels = labels.to(DEVICE)

            wrn_outputs = model(images)
            if len(wrn_outputs) == 4:
                outputs = wrn_outputs[0]
            else:
                outputs = wrn_outputs

            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    return correct, total


def train_knowledge_distillation(teacher, train_dl, epochs, device, learning_rate=0.001, T=2, soft_target_loss_weight=0.25, ce_loss_weight=0.75):
    #instantiate istudent
    student = student_model.Model(num_classes=10).to(device)

    ce_loss = nn.CrossEntropyLoss()
    optimizer = optim.Adam(student.parameters(), lr=learning_rate)
    student_init = copy.deepcopy(student)
    student.to(device)
    teacher.to(device)
    teacher.eval()  # Teacher set to evaluation mode
    student.train() # Student to train mode
    for epoch in range(epochs):
        running_loss = 0.0
        for inputs, labels in train_dl:
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()

            # Forward pass with the teacher model - do not save gradients here as we do not change the teacher's weights
            with torch.no_grad():
                teacher_logits, _, _, _ = teacher(inputs)

            # Forward pass with the student model
            student_logits = student(inputs)
            #Soften the student logits by applying softmax first and log() second
            soft_targets = nn.functional.softmax(teacher_logits / T, dim=-1)
            soft_prob = nn.functional.log_softmax(student_logits / T, dim=-1)

            # Calculate the soft targets loss. Scaled by T**2 as suggested by the authors of the paper "Distilling the knowledge in a neural network"
            soft_targets_loss = torch.sum(soft_targets * (soft_targets.log() - soft_prob)) / soft_prob.size()[0] * (T**2)

            # Calculate the true label loss
            label_loss = ce_loss(student_logits, labels)

            # Weighted sum of the two losses
            loss = soft_target_loss_weight * soft_targets_loss + ce_loss_weight * label_loss

            loss.backward()
            optimizer.step()

            running_loss += loss.item()
        if epoch % 10 == 0:
            print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss / len(train_dl)}")

    return student_init, student


def train_no_cap(model, hp, train_dl, test_dl, optimizer, criterion, scheduler):
    best_test_set_accuracy = 0

    for epoch in range(hp['epochs']):
        model.train()
        for i, data in enumerate(train_dl, 0):
            inputs, labels = data
            inputs = inputs.to(DEVICE)
            labels = labels.to(DEVICE)

            optimizer.zero_grad()

            wrn_outputs = model(inputs)
            if len(wrn_outputs) == 4:
                outputs = wrn_outputs[0]
            else:
                outputs = wrn_outputs
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

        scheduler.step()

        if epoch % 10 == 0 or epoch == hp['epochs'] - 1:
            correct, total = evaluate_on(model, test_dl)
            epoch_accuracy = round(100 * correct / total, 2)
            print(f"Epoch {epoch+1}/{hp['epochs']}: {epoch_accuracy}%")

    return best_test_set_accuracy


def load(hp, model_path, train_dl):
    init_model = model_path / "init_model.pt"
    trained_model = model_path / "trained_model.pt"

    model = WideResNet(
        d=hp["wrn_depth"],
        k=hp["wrn_width"],
        n_classes=10,
        input_features=3,
        output_features=16,
        strides=[1, 1, 2, 2],
    )
    model = ModuleValidator.fix(model)
    ModuleValidator.validate(model, strict=True)
    model_init = copy.deepcopy(model)

    privacy_engine = opacus.PrivacyEngine()
    optimizer = optim.SGD(
        model.parameters(),
        lr=0.1,
        momentum=0.9,
        nesterov=True,
        weight_decay=5e-4
    )
    model, optimizer, train_loader = privacy_engine.make_private_with_epsilon(
        module=model,
        optimizer=optimizer,
        data_loader=train_dl,
        epochs=hp['epochs'],
        target_epsilon=hp['epsilon'],
        target_delta=hp['delta'],
        max_grad_norm=hp['norm'],
    )

    model_init.load_state_dict(torch.load(init_model, weights_only=True))
    model.load_state_dict(torch.load(trained_model, weights_only=True))

    model_init = model_init.to(DEVICE)
    model = model.to(DEVICE)

    adv_points = np.load("data/adv_points.npy")
    adv_labels = np.load("data/adv_labels.npy")
    S = np.load("data/S.npy")

    return model_init, model, adv_points, adv_labels, S



def train_small(hp, train_dl, test_dl):
    model = student_model.Model(num_classes=10).to(DEVICE)
    model = model.to(DEVICE)
    model = ModuleValidator.fix(model)
    ModuleValidator.validate(model, strict=True)

    model_init = copy.deepcopy(model)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    scheduler = MultiStepLR(
        optimizer,
        milestones=[int(i * hp['epochs']) for i in [0.3, 0.6, 0.8]],
        gamma=0.2
    )

    print(f"Training raw (no distill) STUDENT with {hp['epochs']} epochs")

    if hp['epsilon'] is not None:
        privacy_engine = opacus.PrivacyEngine()
        model, optimizer, train_loader = privacy_engine.make_private_with_epsilon(
            module=model,
            optimizer=optimizer,
            data_loader=train_dl,
            epochs=hp['epochs'],
            target_epsilon=hp['epsilon'],
            target_delta=hp['delta'],
            max_grad_norm=hp['norm'],
        )

        print(f"DP epsilon = {hp['epsilon']}, delta = {hp['delta']}")
        print(f"Using sigma={optimizer.noise_multiplier} and C = norm = {hp['norm']}")

        with BatchMemoryManager(
            data_loader=train_loader,
            max_physical_batch_size=2000,  # 1000 ~= 9.4GB vram
            optimizer=optimizer
        ) as memory_safe_data_loader:
            best_test_set_accuracy = train_no_cap(
                model,
                hp,
                memory_safe_data_loader,
                test_dl,
                optimizer,
                criterion,
                scheduler,
            )
    else:
        print("Training without differential privacy")
        best_test_set_accuracy = train_no_cap(
            model,
            hp,
            train_dl,
            test_dl,
            optimizer,
            criterion,
            scheduler,
        )

    return model_init, model

def train(hp, train_dl, test_dl):
    model = WideResNet(
        d=hp["wrn_depth"],
        k=hp["wrn_width"],
        n_classes=10,
        input_features=3,
        output_features=16,
        strides=[1, 1, 2, 2],
    )
    model = model.to(DEVICE)
    model = ModuleValidator.fix(model)
    ModuleValidator.validate(model, strict=True)

    model_init = copy.deepcopy(model)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(
        model.parameters(),
        lr=0.1,
        momentum=0.9,
        nesterov=True,
        weight_decay=5e-4
    )
    scheduler = MultiStepLR(
        optimizer,
        milestones=[int(i * hp['epochs']) for i in [0.3, 0.6, 0.8]],
        gamma=0.2
    )

    print(f"Training with {hp['epochs']} epochs")

    if hp['epsilon'] is not None:
        privacy_engine = opacus.PrivacyEngine()
        model, optimizer, train_loader = privacy_engine.make_private_with_epsilon(
            module=model,
            optimizer=optimizer,
            data_loader=train_dl,
            epochs=hp['epochs'],
            target_epsilon=hp['epsilon'],
            target_delta=hp['delta'],
            max_grad_norm=hp['norm'],
        )

        print(f"DP epsilon = {hp['epsilon']}, delta = {hp['delta']}")
        print(f"Using sigma={optimizer.noise_multiplier} and C = norm = {hp['norm']}")

        with BatchMemoryManager(
            data_loader=train_loader,
            max_physical_batch_size=2000,  # 1000 ~= 9.4GB vram
            optimizer=optimizer
        ) as memory_safe_data_loader:
            best_test_set_accuracy = train_no_cap(
                model,
                hp,
                memory_safe_data_loader,
                test_dl,
                optimizer,
                criterion,
                scheduler,
            )
    else:
        print("Training without differential privacy")
        best_test_set_accuracy = train_no_cap(
            model,
            hp,
            train_dl,
            test_dl,
            optimizer,
            criterion,
            scheduler,
        )

    return model_init, model


def main():
    global DEVICE

    parser = argparse.ArgumentParser(description='WideResNet O1 audit')
    parser.add_argument('--norm', type=float, help='dpsgd norm clip factor', required=True)
    parser.add_argument('--cuda', type=int, help='gpu index', required=False)
    parser.add_argument('--epsilon', type=float, help='dp epsilon', required=False, default=None)
    parser.add_argument('--m', type=int, help='number of target points', required=True)
    parser.add_argument('--k', type=int, help='number of symmetric guesses', required=True)
    parser.add_argument('--epochs', type=int, help='number of epochs', required=True)
    parser.add_argument('--load', type=Path, help='number of epochs', required=False)
    parser.add_argument('--studentraw', action='store_true', help='train a raw student', required=False)
    parser.add_argument('--distill', action='store_true', help='train a raw student', required=False)
    args = parser.parse_args()

    if torch.cuda.is_available() and args.cuda:
        DEVICE = torch.device(f'cuda:{args.cuda}')
    elif torch.cuda.is_available():
        DEVICE = torch.device('cuda:0')
    else:
        DEVICE = torch.device('cpu')

    hp = {
        "target_points": args.m,
        "wrn_depth": 16,
        "wrn_width": 1,
        "epsilon": args.epsilon,
        "delta": 1e-5,
        "norm": args.norm,
        "batch_size": 4096,
        "epochs": args.epochs,
        "k+": args.k,
        "k-": args.k,
        "p_value": 0.05,
    }

    hp['logfile'] = Path('WideResNet_{}_{}_{}_{}s_x{}_{}e_{}d_{}C.txt'.format(
        int(time.time()),
        hp['wrn_depth'],
        hp['wrn_width'],
        hp['batch_size'],
        hp['epochs'],
        hp['epsilon'],
        hp['delta'],
        hp['norm'],
    ))

    if args.load:
        train_dl, test_dl, ____, _, __, ___ = get_dataloaders3(hp['target_points'], hp['batch_size'])
        model_init, model_trained, adv_points, adv_labels, S = load(hp, args.load, train_dl)
        test_dl = None
    else:
        train_dl, test_dl, pure_train_dl, adv_points, adv_labels, S = get_dataloaders3(hp['target_points'], hp['batch_size'])
        if args.studentraw:
            print("=========================")
            print("Training a raw student model")
            print("=========================")
            model_init, model_trained = train_small(hp, train_dl, test_dl)
        elif args.distill:
            print("=========================")
            print("Training a distilled student model")
            print("=========================")
            teacher_init, teacher_trained = train(hp, train_dl, test_dl)
            model_init, model_trained = train_knowledge_distillation(
                teacher=teacher_trained,
                train_dl=train_dl,
                epochs=100,
                device=DEVICE,
                learning_rate=0.001,
                T=2,
                soft_target_loss_weight=0.25,
                ce_loss_weight=0.75,
            )
        else:
            print("=========================")
            print("Training teacher model")
            print("=========================")
            model_init, model_trained = train(hp, train_dl, test_dl)

        np.save("data/adv_points", adv_points)
        np.save("data/adv_labels", adv_labels)
        np.save("data/S", S)
        torch.save(model_init.state_dict(), "data/init_model.pt")
        torch.save(model_trained.state_dict(), "data/trained_model.pt")

    scores = list()
    criterion = nn.CrossEntropyLoss()
    with torch.no_grad():
        model_init.eval()
        x_m = torch.from_numpy(adv_points).to(DEVICE)
        y_m = torch.from_numpy(adv_labels).long().to(DEVICE)

        for i in range(len(x_m)):
            x_point = x_m[i].unsqueeze(0).to(DEVICE)
            y_point = y_m[i].unsqueeze(0).to(DEVICE)
            is_in = S[i]

            wrn_outputs = model_init(x_point)
            outputs = wrn_outputs[0] if len(wrn_outputs) == 4 else wrn_outputs
            init_loss = criterion(outputs, y_point)

            wrn_outputs = model_trained(x_point)
            outputs = wrn_outputs[0] if len(wrn_outputs) == 4 else wrn_outputs
            trained_loss = criterion(outputs, y_point)

            scores.append(((init_loss - trained_loss).item(), is_in))

    scores = sorted(scores, key=lambda x: x[0])
    scores = np.array([x[1] for x in scores])

    print(scores[:10])

    audits = (0, 0, 0, 0)
    for k in [10, 20, 50, 100, 200, 300, 500, 800, 1000, 1200, 1400, 1600, 1800, 1900, 2000, 2100, 2200, 2300, 2400, 2500]:
        correct = np.sum(~scores[:k]) + np.sum(scores[-k:])
        total = len(scores)

        eps_lb = get_eps_audit(
            hp['target_points'],
            2*k,
            correct,
            hp['delta'],
            hp['p_value']
        )

        if eps_lb > audits[0]:
            audits = (eps_lb, k, correct, total)

    print(f"Audit total: {audits[2]}/{2*audits[1]}/{audits[3]}")
    print(f"p[ε < {audits[0]}] < {hp['p_value']} for true epsilon {hp['epsilon']}")

    if test_dl is not None:
        correct, total = evaluate_on(model_init, test_dl)
        print(f"Init model accuracy: {correct}/{total} = {round(correct/total*100, 2)}")
        correct, total = evaluate_on(model_trained, test_dl)
        print(f"Done model accuracy: {correct}/{total} = {round(correct/total*100, 2)}")


if __name__ == '__main__':
    main()