def test_extraction():
"""End-to-end test of a model extraction attack"""
# Create a query function for a target PyTorch Lightning model
model = train_four_layer_mnist_victim(gpus=torch.cuda.device_count())
def query_mnist(input_data):
# PrivacyRaven provides built-in query functions
return get_target(model, input_data, (1, 28, 28, 1))
# Obtain seed (or public) data to be used in extraction
emnist_train, emnist_test = get_emnist_data()
# Run a model extraction attack
attack = ModelExtractionAttack(
query=query_mnist,
query_limit=100,
victim_input_shape=(1, 28, 28, 1), # EMNIST data point shape
victim_output_targets=10,
substitute_input_shape=(3, 1, 28, 28),
synthesizer="copycat",
substitute_model_arch=FourLayerClassifier, # 28*28: image size
substitute_input_size=784,
seed_data_train=emnist_train,
seed_data_test=emnist_test,
gpus=0,
)
Users should refer to https://pytorch-lightning-bolts.readthedocs.io/en/latest/callbacks.html to construct
Pytorch Lightning Bolt callbacks.
"""
import privacyraven as pr
from privacyraven.utils.data import get_emnist_data
from pl_bolts.callbacks import PrintTableMetricsCallback
from privacyraven.extraction.core import ModelExtractionAttack
from privacyraven.utils.query import get_target
from privacyraven.models.victim import train_four_layer_mnist_victim
from privacyraven.models.four_layer import FourLayerClassifier
from pytorch_lightning.callbacks import Callback
# Trains a 4-layer fully connected neural network on MNIST data using all of the GPUs
# available to the user, or CPU if no GPUs are available (torch.cuda.device_count handles this).
model = train_four_layer_mnist_victim(gpus=torch.cuda.device_count())
# Create a query function for a target PyTorch Lightning model
def query_mnist(input_data):
# PrivacyRaven provides built-in query functions
return get_target(model, input_data, (1, 28, 28, 1))
emnist_train, emnist_test = get_emnist_data()
# Below is a user-defined callback that inherits from the Pytorch's Lightning Bolt Callback class.
# All it does is print "End of epoch" at the end of a training epoch.
class CustomCallback(Callback):
This model extraction attack steals a model trained on MNIST by
using the copycat synthesizer and the EMNIST dataset to train a
FourLayerClassifier substitute. A single GPU is assumed.
"""
import privacyraven as pr
from privacyraven.utils.data import get_emnist_data
from privacyraven.extraction.core import ModelExtractionAttack
from privacyraven.utils.query import get_target
from privacyraven.models.victim import train_four_layer_mnist_victim
from privacyraven.models.four_layer import FourLayerClassifier
# Trains a 4-layer fully connected neural network on MNIST data using the user's GPUs. See
# src/privacyraven/models/victims.py for a full set of supported parameters.
model = train_four_layer_mnist_victim(gpus=1)
# Create a query function for a target PyTorch Lightning model
def query_mnist(input_data):
# PrivacyRaven provides built-in query functions
return get_target(model, input_data, (1, 28, 28, 1))
# Obtain seed (or public) data to be used in extraction
emnist_train, emnist_test = get_emnist_data()
# Run a model extraction attack
attack = ModelExtractionAttack(
query_mnist,
200, # Less than the number of MNIST data points: 60000