# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Membership Inference Attack (MIA)
#
# We demonstate *transfer-inherit shadow learning* which is membership inference attack for collaborative inference like SplitNN.
#
# reference  
# https://ieeexplore.ieee.org/document/9302683

from attacksplitnn.splitnn import Client, Server, SplitNN
from attacksplitnn.attack import TransferInherit
from attacksplitnn.utils import DataSet

# +
import numpy as np
import pandas as pd
import random
import matplotlib.pyplot as plt

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data.dataset import Dataset
from torch.utils.data import DataLoader

from sklearn.svm import SVC
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, roc_auc_score
# -

# ## Set paramters

torch.random.manual_seed(42)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)

# ## Load and preprocess data

trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True)
testset = torchvision.datasets.MNIST(root='./data', train=False, download=True)

# +
X_train = np.array(trainset.data)
y_train = np.array(trainset.targets)

X_test = np.array(testset.data)
y_test = np.array(testset.targets)

print(X_train.shape, y_train.shape)
print(X_test.shape, y_test.shape)

# +
victim_idx = random.sample(range(X_train.shape[0]), k=4000)
attack_idx = random.sample(range(X_test.shape[0]), k=3000)
shadow_idx = attack_idx[:1000]
eval_idx = attack_idx[1000:]

X_victim = X_train[victim_idx]
y_victim = y_train[victim_idx]

X_shadow = X_test[shadow_idx]
y_shadow = y_test[shadow_idx]

X_eval = X_test[eval_idx]
y_eval = y_test[eval_idx]

print(X_victim.shape, y_victim.shape)
print(X_shadow.shape, y_shadow.shape)
print(X_eval.shape, y_eval.shape)

# +
transform = transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])

victimset = DataSet(X_victim, y_victim, transform=transform)
victimloader = torch.utils.data.DataLoader(victimset, batch_size=64, shuffle=True)

shadowset = DataSet(X_shadow, y_shadow, transform=transform)
shadowloader = torch.utils.data.DataLoader(shadowset, batch_size=64, shuffle=True)

valset = DataSet(X_eval, y_eval, transform=transform)
valloader = torch.utils.data.DataLoader(valset, batch_size=64, shuffle=True)


# -

# ## Train victim SplitNN

class FirstNet(nn.Module):
    def __init__(self):
        super(FirstNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=64,
                               kernel_size=3, padding=1, stride=1)
        self.bn1 = nn.BatchNorm2d(64)
        self.conv2 = nn.Conv2d(64, 128, 3, padding=1) 
        self.bn2 = nn.BatchNorm2d(128)

    def forward(self, x):
        # 3ch > 64ch, shape 32 x 32 > 16 x 16
        x = self.conv1(x) # [64,32,32]
        x = self.bn1(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2, 2) # [64,16,16]
        
        # 64ch > 128ch, shape 16 x 16 > 8 x 8
        x = self.conv2(x) # [128,16,16]
        x = self.bn2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2, 2) # [128,8,8]
        return x


# CNNを実装する
class SecondNet(nn.Module):
    def __init__(self):
        super(SecondNet, self).__init__()
        self.conv3 = nn.Conv2d(128, 256, 3, padding=1)
        self.bn3 = nn.BatchNorm2d(256)
        self.conv4 = nn.Conv2d(256, 512, 3, padding=1)
        self.bn4 = nn.BatchNorm2d(512)
        self.L1 = nn.Linear(512, 10) # 10クラス分類

    def forward(self, x):
        # 128ch > 256ch, shape 8 x 8 > 4 x 4
        x = self.conv3(x) # [256,8,8]
        x = self.bn3(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2, 2) # [256,4,4]   

        # 256ch > 512ch, shape 4 x 4 > 2 x 2
        x = self.conv4(x) # [512,4,4]
        x = self.bn4(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2, 2) # [512,2,2]
        # 全結合層
        x = x.view(-1, 512)
        x = self.L1(x)
        #x = F.softmax(x, dim=0)
        return x



# +
model_1 = FirstNet()
model_1 = model_1.to(device)

model_2 = SecondNet()
model_2 = model_2.to(device)

opt_1 = optim.Adam(model_1.parameters(), lr=1e-3)
opt_2 = optim.Adam(model_2.parameters(), lr=1e-3)

criterion = nn.CrossEntropyLoss()


# -

def accuracy(label, output):
    pred = output.argmax(dim=1, keepdim=True)
    return pred.eq(label.view_as(pred)).sum().item() / pred.shape[0]


# +
client = Client(model_1)
server = Server(model_2)

splitnn = SplitNN(client, server, opt_1, opt_2)
# -

splitnn.train()
for epoch in range(3):
    epoch_loss = 0
    epoch_outputs = []
    epoch_labels = []
    for i, data in enumerate(victimloader):
        splitnn.zero_grads()
        inputs, labels = data
        inputs = inputs.to(device)
        labels = labels.to(device)
        
        outputs = splitnn(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        epoch_loss += loss.item() / len(victimloader.dataset)
        
        epoch_outputs.append(outputs)
        epoch_labels.append(labels)
        
        splitnn.backward()
        splitnn.step()
        
    print(epoch_loss, accuracy(torch.cat(epoch_labels),
                                torch.cat(epoch_outputs)))

# ## Train shadow model
#
# The server wants to execute a membership inference attack against the client, so the server needs the shadow model to mimic the client. The server uses pre-trained alexnet as a shadow model and concatenates it with the server-side model to create the system that imitates the victim SplitNN. Then, the server can create the membership dataset and train the attacker classifier. In this notebook, we use SVM as an attacker model. 

# +
model_transfer = torch.hub.load('pytorch/vision:v0.9.0', 'alexnet', pretrained=True)
model_1_shadow = model_transfer.features
model_1_shadow[0] = nn.Conv2d(1, 64, kernel_size=(1, 1), stride=(1, 1))
model_1_shadow = nn.Sequential(model_1_shadow,
                               nn.Conv2d(256, 128, kernel_size=(1, 1),
                                         stride=(2, 2), padding=(6, 6)))
opt_3 = optim.Adam(model_1_shadow.parameters(), lr=1e-3)

shadow_client = Client(model_1_shadow)
# -

attacker_clf = SVC(probability=True)

smia = TransferInherit(splitnn, shadow_client, opt_3, attacker_clf, device=device)

smia.fit(shadowloader, valloader,
         5, criterion, shadow_metric=accuracy,
         attack_dataset_split=0.33, random_state=42)

# ## The performance of Membership Inference Attack

# +
pred_proba = smia.attack(smia.attacker_X_train)[:,1]
print("train AUC", roc_auc_score(smia.attacker_y_train, pred_proba))

pred_proba = smia.attack(smia.attacker_X_test)[:,1]
print("test AUC", roc_auc_score(smia.attacker_y_test, pred_proba))
# -



StarcoderdataJupyterScriptsDedupFiltered