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