Certification of Robustness using Zonotopes with DeepZ¶

# train the model normally

def standard_train(model, opt, criterion, x, y, bsize=32, epochs=5):
    num_of_batches = int(len(x) / bsize)
    for epoch in range(epochs):
        x, y = shuffle(x, y)
        loss_list = []
        for bnum in range(num_of_batches):
            x_batch = np.copy(x[bnum * bsize:(bnum + 1) * bsize])
            y_batch = np.copy(y[bnum * bsize:(bnum + 1) * bsize])

            x_batch = torch.from_numpy(x_batch).float().to(device)
            y_batch = torch.from_numpy(y_batch).type(torch.LongTensor).to(device)

            # zero the parameter gradients
            opt.zero_grad()
            outputs = model(x_batch)
            loss = criterion(outputs, y_batch)
            loss_list.append(loss.data)
            loss.backward()
            opt.step()
        print('End of epoch {} loss {}'.format(epoch, np.mean(loss_list)))
    return model

model = standard_train(model=model,
                       opt=opt,
                       criterion=criterion,
                       x=x_train,
                       y=y_train)

# lets now get the predicions for the MNIST test set and see how well our model is doing.
with torch.no_grad():
    test_preds = model(torch.from_numpy(x_test).float().to(device))

test_preds = np.argmax(test_preds.cpu().detach().numpy(), axis=1)
print('Test acc: ', np.mean(test_preds == y_test) * 100)

# But how robust are these predictions? 
# We can now examine this neural network's certified robustness. 
# We pass it into PytorchDeepZ. We will get a print out showing which 
# neural network layers have been registered. There will also be a 
# warning to tell us that PytorchDeepZ currently infers a reshape when 
# a neural network goes from using convolutional to dense layers. 
# This will cover the majority of use cases, however, if not then the 
# certification layers in art.estimators.certification.deepz.deep_z.py 
# can be used to directly build a certified model structure.

zonotope_model = deep_z.PytorchDeepZ(model=model, 
                                     clip_values=(0, 1), 
                                     loss=nn.CrossEntropyLoss(), 
                                     input_shape=(1, 28, 28), 
                                     nb_classes=10)

# Lets now see how robust our model is!
# First we need to define what bound we need to check. 
# Here let's check for L infinity robustness with small bound of 0.05

bound = 0.05
num_certified = 0
num_correct = 0

# lets now loop over the data to check its certified robustness:
# we need to consider a single sample at a time as due to memory and compute footprints batching is not supported.
# In this demo we will look at the first 50 samples of the MNIST test data.
original_x = np.copy(x_test)
for i, (sample, pred, label) in enumerate(zip(x_test[:50], test_preds[:50], y_test[:50])):
    
    # we make the matrix representing the allowable perturbations. 
    # we have 28*28 features and each one can be manipulated independently requiring a different row.
    # hence a 784*784 matrix.
    eps_bound = np.eye(784) * bound
    
    # we then need to adjust the raw data with the eps bounds to take into account
    # the allowable range of 0 - 1 for pixel data.
    # We provide a simple function to do this preprocessing for image data.
    # However if your use case is not supported then a custom pre-processor function will need to be written.
    sample, eps_bound = zonotope_model.pre_process(cent=sample, 
                                                   eps=eps_bound)
    sample = np.expand_dims(sample, axis=0)

    # We pass the data sample and the eps bound to the certifier along with the prediction that was made
    # for the datapoint. 
    # A boolean is returned signifying if it can have its class changed under the given bound.
    is_certified = zonotope_model.certify(cent=sample,
                                          eps=eps_bound,
                                          prediction=pred)
    
    if pred == label:
        num_correct +=1
        if is_certified:
            num_certified +=1 
    
    print('Classified Correct {}/{} and also certified {}/{}'.format(num_correct, i+1, num_certified, i+1))

# we can then compare this to the empirical PGD performance

from art.estimators.classification import PyTorchClassifier
from art.attacks.evasion.projected_gradient_descent.projected_gradient_descent import ProjectedGradientDescent

classifier = PyTorchClassifier(
    model=model,
    clip_values=(0.0, 1.0),
    loss=criterion,
    optimizer=opt,
    input_shape=(1, 28, 28),
    nb_classes=10,
)

attack = ProjectedGradientDescent(classifier, eps=0.05, eps_step=0.01, verbose=False)
x_train_adv = attack.generate(x_test[:50].astype('float32'))
y_adv_pred = classifier.predict(torch.from_numpy(x_train_adv).float().to(device))
y_adv_pred = np.argmax(y_adv_pred, axis=1)
print('Test acc: ', np.mean(y_adv_pred == y_test[:50]) * 100)