#!/usr/bin/env python
# -*- coding: utf-8 -*-
# References: J. Buckman, et al., "Thermometer Encoding: One hot way to resist adversarial examples," in ICLR, 2018.
# Reference Implementation (TensorFlow): https://github.com/anishathalye/obfuscated-gradients/tree/master/thermometer
# **************************************
# @Time : 2018/11/22 18:30
# @Author : Jiaxu Zou
# @Lab : nesa.zju.edu.cn
# @File : TE.py
# **************************************
import os
import numpy as np
import torch
import torch.nn.functional as F
import torch.optim as optim
from Defenses.DefenseMethods.defenses import Defense
from RawModels.ResNet import adjust_learning_rate
def one_hot_encoding(samples=None, level=None, device=None):
"""
the help function to encode the samples using the one-hot encoding schema
:param samples:
:param level:
:param device:
:return:
"""
assert level is not None and isinstance(level, int), 'level should specified as an integer'
assert torch.is_tensor(samples), "input samples must be a PyTorch tensor"
if len(samples.shape) >= 4 and (samples.shape[1] == 1 or samples.shape[1] == 3):
samples = samples.permute(0, 2, 3, 1)
# inserting the last position for samples (handle the upper bound by multiplying 0.9999)
discretized_samples = torch.unsqueeze(input=(0.99999 * samples * level).long().to(device), dim=4)
# make the last dim be the level number
shape = discretized_samples.shape
# convert to one_hot encoding
one_hot_samples = torch.zeros([shape[0], shape[1], shape[2], shape[3], level]).to(device).scatter_(-1, discretized_samples, 1)
one_hot_samples = one_hot_samples.float()
return one_hot_samples
def thermometer_encoding(samples=None, level=None, device=None):
"""
the help function to encode the samples using the thermometer encoding schema
:param samples:
:param level:
:param device:
:return:
"""
assert level is not None and isinstance(level, int), 'level should specified as an integer'
assert torch.is_tensor(samples), "input samples must be a PyTorch tensor"
if len(samples.shape) >= 4 and (samples.shape[1] == 1 or samples.shape[1] == 3):
samples = samples.permute(0, 2, 3, 1)
# convert one hot encoding to thermometer encoding
one_hot_samples = one_hot_encoding(samples=samples, level=level, device=device)
therm_samples = torch.cumsum(one_hot_samples, dim=-1)
# the returned samples is a type of numpy data with shape [BatchSize * (Channel * Level) * Weight* Height]
shape = samples.shape
therm_samples_numpy = torch.reshape(therm_samples, (shape[0], shape[1], shape[2], shape[3] * level))
therm_samples_numpy = therm_samples_numpy.permute(0, 3, 1, 2).cpu().numpy()
return therm_samples_numpy
class TEDefense(Defense):
def __init__(self, model=None, defense_name=None, dataset=None, training_parameters=None, device=None, **kwargs):
"""
:param model:
:param defense_name:
:param dataset:
:param training_parameters:
:param device:
:param kwargs:
"""
super(TEDefense, self).__init__(model=model, defense_name=defense_name)
self.model = model
self.defense_name = defense_name
self.device = device
self.Dataset = dataset.upper()
assert self.Dataset in ['MNIST', 'CIFAR10'], "The data set must be MNIST or CIFAR10"
# make sure to parse the parameters for the defense
assert self._parsing_parameters(**kwargs)
# get the training_parameters, the same as the settings of RawModels
self.num_epochs = training_parameters['num_epochs']
self.batch_size = training_parameters['batch_size']
# prepare the optimizers
if self.Dataset == 'MNIST':
self.optimizer = optim.SGD(self.model.parameters(), lr=training_parameters['learning_rate'],
momentum=training_parameters['momentum'], weight_decay=training_parameters['decay'], nesterov=True)
else:
self.optimizer = optim.Adam(self.model.parameters(), lr=training_parameters['lr'])
def _parsing_parameters(self, **kwargs):
assert kwargs is not None, "the parameters should be specified"
print("\nUser configurations for the {} defense".format(self.defense_name))
for key in kwargs:
print('\t{} = {}'.format(key, kwargs[key]))
self.level = kwargs['level']
self.steps = kwargs['steps']
self.attack_eps = kwargs['attack_eps']
self.attack_step_size = kwargs['attack_step_size']
return True
def lspga_generation(self, samples=None, ys=None, noise_init=True):
"""
one type of white-box attacks on discretized inputs (thermometer encoding) -- Logit-Space Projected Gradient Ascent (LS-PGA)
the detailed pseudo-code for LS-PGA attack is described in Algorithm 3 of the referenced paper
:param samples:
:param ys:
:param noise_init:
:return:
"""
# STEP 1: sub-routine for getting an \epsilon-discretized masked of an image
lowest = torch.clamp(samples - self.attack_eps, 0.0, 1.0)
highest = torch.clamp(samples + self.attack_eps, 0.0, 1.0)
# get the masking of intervals between lowest and highest
masked_intervals = 0.0
for alpha in np.linspace(0., 1., self.level):
single_one_hot = one_hot_encoding(samples=alpha * lowest + (1. - alpha) * highest, level=self.level, device=self.device)
masked_intervals += single_one_hot
masked = (masked_intervals > 0.0).float()
shape = masked.shape
# STEP 2: main function for generating adversarial examples using LS-PGA
# init each of logits randomly with values sampled from a standard normal distribution.
if noise_init is True:
us_numpy = torch.randn(shape).cpu().numpy()
else:
us_numpy = torch.zeros_like(masked).cpu().numpy()
# generating
inv_temp = 1.0
sigma_rate = 1.2
self.model.eval()
for i in range(self.steps):
us_logits = torch.from_numpy(us_numpy).to(self.device).float()
us_logits.requires_grad = True
# if not masked ( equal 0), turn it to be -inf (-99999)
# then embedding the logits using softmax function with temperature to
us_probs = F.softmax(inv_temp * (us_logits * masked - 999999.0 * (1. - masked)), dim=-1)
# apply the cumulative sum function and reshape to get the distribution embedding
thermometer_probs = torch.cumsum(us_probs, dim=-1)
thermometer_probs = torch.reshape(thermometer_probs, (shape[0], shape[1], shape[2], shape[3] * self.level))
# convert the channel back to the second position
thermometer_probs = thermometer_probs.permute(0, 3, 1, 2)
logits = self.model(thermometer_probs)
if ys is None and i == 0:
ys = torch.argmax(logits, dim=1)
loss = F.cross_entropy(logits, ys)
gradients = torch.autograd.grad(loss, us_logits)[0]
signed_gradient = torch.sign(gradients).cpu().numpy()
us_numpy += self.attack_step_size * signed_gradient
inv_temp *= sigma_rate # anneal the temperature via exponential decay with rate sigma
us_logits = torch.from_numpy(us_numpy).to(self.device).float()
logits_results = us_logits * masked - 999999.0 * (1. - masked)
logits_final = torch.argmax(logits_results, dim=-1, keepdim=True)
one_hot_adv_samples = torch.zeros([shape[0], shape[1], shape[2], shape[3], self.level]).to(self.device).scatter_(-1, logits_final, 1)
one_hot_adv_samples = one_hot_adv_samples.float()
# the returned samples is a type of numpy dataset
therm_adv_samples = torch.cumsum(one_hot_adv_samples, dim=-1)
final_adv_samples = torch.reshape(therm_adv_samples, (shape[0], shape[1], shape[2], shape[3] * self.level))
final_adv_samples_numpy = final_adv_samples.permute(0, 3, 1, 2).cpu().numpy()
return final_adv_samples_numpy
def train_one_epoch_with_adv_lspga(self, train_loader=None, epoch=None, weight_regular=None):
"""
:param train_loader:
:param epoch:
:return:
"""
for index, (images, labels) in enumerate(train_loader):
nat_images_numpy = thermometer_encoding(samples=images.to(self.device), level=self.level, device=self.device)
nat_labels = labels.to(self.device)
# prepare for LSPGA perturbation
self.model.eval()
adv_images_numpy = self.lspga_generation(samples=images.to(self.device))
# concatenate the nature samples and adversarial examples
batch_images_numpy = np.concatenate((nat_images_numpy, adv_images_numpy), axis=0)
batch_images = torch.from_numpy(batch_images_numpy).to(self.device)
# concatenate the true labels
batch_labels = torch.cat((nat_labels, nat_labels), dim=0)
# set the model in the training mode
self.model.train()
# forward the nn
logits = self.model(batch_images)
loss = F.cross_entropy(logits, batch_labels)
# backward
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
print('\rTrain Epoch{:>3}: [batch:{:>4}/{:>4}] \tLoss={:.4f} ===> '.format(epoch, index, len(train_loader), loss), end=' ')
def thermometer_validation_evaluation(self, validation_loader, device):
"""
validation evaluation with slight modification for thermometer encoded input samples
:param validation_loader:
:param device:
:return:
"""
self.model.eval()
total = 0.0
correct = 0.0
with torch.no_grad():
for index, (inputs, labels) in enumerate(validation_loader):
therm_inputs = thermometer_encoding(samples=inputs.to(self.device), level=self.level, device=device)
therm_inputs = torch.from_numpy(therm_inputs).to(self.device)
labels = labels.to(device)
outputs = self.model(therm_inputs)
_, predicted = torch.max(outputs.data, 1)
total = total + labels.size(0)
correct = correct + (predicted == labels).sum().item()
ratio = correct / total
print('validation set accuracy is ', ratio)
return ratio
def defense(self, train_loader=None, validation_loader=None):
best_val_acc = None
for epoch in range(self.num_epochs):
# training the model with nature examples and corresponding adversarial examples
self.train_one_epoch_with_adv_lspga(train_loader=train_loader, epoch=epoch, weight_regular=1e-4)
val_acc = self.thermometer_validation_evaluation(validation_loader=validation_loader, device=self.device)
# adjust the learning rate for cifar10 training
if self.Dataset == 'CIFAR10':
adjust_learning_rate(optimizer=self.optimizer, epoch=epoch)
# save the retained defense-enhanced model
assert os.path.exists('../DefenseEnhancedModels/{}'.format(self.defense_name))
defense_enhanced_saver = '../DefenseEnhancedModels/{}/{}_{}_enhanced.pt'.format(self.defense_name, self.Dataset, self.defense_name)
if not best_val_acc or round(val_acc, 4) >= round(best_val_acc, 4):
if best_val_acc is not None:
os.remove(defense_enhanced_saver)
best_val_acc = val_acc
self.model.save(name=defense_enhanced_saver)
else:
print('Train Epoch{:>3}: validation dataset accuracy did not improve from {:.4f}\n'.format(epoch, best_val_acc))
