def test_extraction():
try:
print("Creating victim model")
model = train_mnist_victim(gpus=0)
def query_mnist(input_data):
return get_target(model, input_data)
print("Downloading EMNIST data")
emnist_train, emnist_test = get_emnist_data()
print("Launching model extraction attack")
# A single GPU is assumed
attack = ModelExtractionAttack(
query=query_mnist,
query_limit=100,
victim_input_shape=(1, 28, 28, 1),
victim_output_targets=10,
substitute_input_shape=(1, 3, 28, 28),
synthesizer="copycat",
substitute_model_arch=ImagenetTransferLearning,
substitute_input_size=1000,
seed_data_train=emnist_train,
seed_data_test=emnist_test,
gpus=0,
)
print(attack)
except Exception:
pytest.fail("Unexpected Error")
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,
)
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):
def on_epoch_end(self, trainer, pl_module):
print('End of epoch')
# Runs a Model Extraction Attack with the user-defined CustomCallback specified as an argument.
# Note that parentheses are needed while passing in the callback, since
# Pytorch Lightning bolt callbacks are classes that need to be instantiated.
attack = ModelExtractionAttack(query=query_mnist,