Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
177 lines (131 sloc) 5.23 KB
"""
Example membership inference attack against a deep net classifier on the CIFAR10 dataset
"""
try:
import keras
except ImportError as e:
raise ImportError(
"You need to have keras installed for this example: pip install keras"
)
import numpy as np
from absl import app
from absl import flags
from keras import layers
from keras.datasets import cifar10
from sklearn.model_selection import train_test_split
from mia.estimators import ShadowModelBundle, AttackModelBundle, prepare_attack_data
NUM_CLASSES = 10
WIDTH = 32
HEIGHT = 32
CHANNELS = 3
SHADOW_DATASET_SIZE = 4000
ATTACK_TEST_DATASET_SIZE = 4000
FLAGS = flags.FLAGS
flags.DEFINE_integer(
"target_epochs", 12, "Number of epochs to train target and shadow models."
)
flags.DEFINE_integer("attack_epochs", 12, "Number of epochs to train attack models.")
flags.DEFINE_integer("num_shadows", 3, "Number of epochs to train attack models.")
def get_data():
"""Prepare CIFAR10 data."""
(X_train, y_train), (X_test, y_test) = cifar10.load_data()
y_train = keras.utils.to_categorical(y_train)
y_test = keras.utils.to_categorical(y_test)
X_train = X_train.astype("float32")
X_test = X_test.astype("float32")
y_train = y_train.astype("float32")
y_test = y_test.astype("float32")
X_train /= 255
X_test /= 255
return (X_train, y_train), (X_test, y_test)
def target_model_fn():
"""The architecture of the target (victim) model.
The attack is white-box, hence the attacker is assumed to know this architecture too."""
model = keras.models.Sequential()
model.add(
layers.Conv2D(
32,
(3, 3),
activation="relu",
padding="same",
input_shape=(WIDTH, HEIGHT, CHANNELS),
)
)
model.add(layers.Conv2D(32, (3, 3), activation="relu"))
model.add(layers.MaxPooling2D(pool_size=(2, 2)))
model.add(layers.Dropout(0.25))
model.add(layers.Conv2D(64, (3, 3), activation="relu", padding="same"))
model.add(layers.Conv2D(64, (3, 3), activation="relu"))
model.add(layers.MaxPooling2D(pool_size=(2, 2)))
model.add(layers.Dropout(0.25))
model.add(layers.Flatten())
model.add(layers.Dense(512, activation="relu"))
model.add(layers.Dropout(0.5))
model.add(layers.Dense(NUM_CLASSES, activation="softmax"))
model.compile("adam", loss="categorical_crossentropy", metrics=["accuracy"])
return model
def attack_model_fn():
"""Attack model that takes target model predictions and predicts membership.
Following the original paper, this attack model is specific to the class of the input.
AttachModelBundle creates multiple instances of this model for each class.
"""
model = keras.models.Sequential()
model.add(layers.Dense(128, activation="relu", input_shape=(NUM_CLASSES,)))
model.add(layers.Dropout(0.3, noise_shape=None, seed=None))
model.add(layers.Dense(64, activation="relu"))
model.add(layers.Dropout(0.2, noise_shape=None, seed=None))
model.add(layers.Dense(64, activation="relu"))
model.add(layers.Dense(1, activation="sigmoid"))
model.compile("adam", loss="binary_crossentropy", metrics=["accuracy"])
return model
def demo(argv):
del argv # Unused.
(X_train, y_train), (X_test, y_test) = get_data()
# Train the target model.
print("Training the target model...")
target_model = target_model_fn()
target_model.fit(
X_train, y_train, epochs=FLAGS.target_epochs, validation_split=0.1, verbose=True
)
# Train the shadow models.
smb = ShadowModelBundle(
target_model_fn,
shadow_dataset_size=SHADOW_DATASET_SIZE,
num_models=FLAGS.num_shadows,
)
# We assume that attacker's data were not seen in target's training.
attacker_X_train, attacker_X_test, attacker_y_train, attacker_y_test = train_test_split(
X_test, y_test, test_size=0.1
)
print(attacker_X_train.shape, attacker_X_test.shape)
print("Training the shadow models...")
X_shadow, y_shadow = smb.fit_transform(
attacker_X_train,
attacker_y_train,
fit_kwargs=dict(
epochs=FLAGS.target_epochs,
verbose=True,
validation_data=(attacker_X_test, attacker_y_test),
),
)
# ShadowModelBundle returns data in the format suitable for the AttackModelBundle.
amb = AttackModelBundle(attack_model_fn, num_classes=NUM_CLASSES)
# Fit the attack models.
print("Training the attack models...")
amb.fit(
X_shadow, y_shadow, fit_kwargs=dict(epochs=FLAGS.attack_epochs, verbose=True)
)
# Test the success of the attack.
# Prepare examples that were in the training, and out of the training.
data_in = X_train[:ATTACK_TEST_DATASET_SIZE], y_train[:ATTACK_TEST_DATASET_SIZE]
data_out = X_test[:ATTACK_TEST_DATASET_SIZE], y_test[:ATTACK_TEST_DATASET_SIZE]
# Compile them into the expected format for the AttackModelBundle.
attack_test_data, real_membership_labels = prepare_attack_data(
target_model, data_in, data_out
)
# Compute the attack accuracy.
attack_guesses = amb.predict(attack_test_data)
attack_accuracy = np.mean(attack_guesses == real_membership_labels)
print(attack_accuracy)
if __name__ == "__main__":
app.run(demo)