def defense(self, train_loader=None, validation_loader=None):
# train_loader：训练集
# validation_loader：验证集
# best_val_acc：验证集上的最佳分类精度
best_val_acc = None
# 进行epoch次训练
for epoch in range(self.num_epochs):
self.train_one_epoch_with_adv_and_nat(train_loader=train_loader,
epoch=epoch)
# 计算每一次训练后在验证集上的分类精度
val_acc = validation_evaluation(
model=self.model,
validation_loader=validation_loader,
device=self.device)
# 如果是CIFAR10数据集，则调整模型参数
if self.Dataset == 'CIFAR10':
adjust_learning_rate(epoch=epoch, optimizer=self.optimizer)
# 将最佳模型参数保存到DefenseEnhancedModels/NAT/CIFAR10_NAT_enhanced.pt中或MNIST
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))
def defense(self, train_loader=None, validation_loader=None):
# train_loader：训练集
# validation_loader：验证集
# best_val_acc：验证集最佳分类精度
best_val_acc = None
# 进行num_epochs次PGD对抗训练
for epoch in range(self.num_epochs):
# 进行一次完整PGD对抗训练
self.train_one_epoch_with_pgd_and_nat(train_loader=train_loader, epoch=epoch)
# val_acc：对验证集进行评估的分类精度
val_acc = validation_evaluation(model=self.model, validation_loader=validation_loader, device=self.device)
# 进行学习率调整
if self.Dataset == 'CIFAR10':
adjust_CIFAR10_learning_rate(epoch=epoch, optimizer=self.optimizer)
else:
adjust_MNIST_learning_rate(epoch=epoch, optimizer=self.optimizer)
assert os.path.exists('../DefenseEnhancedModels/{}'.format(self.defense_name))
# defense_enhanced_saver：对抗训练网络参数的存放位置为DefenseEnhancedModels/PAT/CIFAR10_PAT_enhanced.pt中
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))
def defense(self, pre_trained_models=None, train_loader=None, validation_loader=None):
# pre_trained_models：预训练模型
# train_loader：训练集
# validation_loader：验证集
# best_val_acc：最佳测试集分类精度
best_val_acc = None
# 进行num_epochs次集成对抗训练
for epoch in range(self.num_epochs):
# training the model with natural examples and corresponding adversarial examples from external models
# 一次完整的集成对抗训练
self.train_one_epoch_with_adv_from_external_models(pre_trained_models=pre_trained_models, train_loader=train_loader, epoch=epoch)
# 计算训练后的模型对验证集的分类精度
val_acc = validation_evaluation(model=self.model, validation_loader=validation_loader, device=self.device)
# 如果数据集为CIFAR10则进行学习率的调整
if self.Dataset == 'CIFAR10':
adjust_learning_rate(epoch=epoch, optimizer=self.optimizer_adv)
assert os.path.exists('../DefenseEnhancedModels/{}'.format(self.defense_name))
# defense_enhanced_saver：集成对抗训练防御模型参数的保存位置在DefenseEnhancedModels/NEAT/MNIST_NEAT_enhanced.pt中或者CIFAR10
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))
def main(args):
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_index
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.manual_seed(args.seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(args.seed)
np.random.seed(args.seed)
random.seed(args.seed)
# 下载mnist训练集，划分为训练集和测试集（然后进行分组），保存到MNIST文件夹下面
train_loader, valid_loader = get_mnist_train_validate_loader(dir_name='../data/MNIST', batch_size=MNIST_Training_Parameters['batch_size'],valid_size=0.1, shuffle=True)
# 下载minst测试集（然后进行分组），保存到MNIST文件夹下面
test_loader = get_mnist_test_loader(dir_name='../data/MNIST', batch_size=MNIST_Training_Parameters['batch_size'])
# 设置模型
# **************引入的模型名称**************
mnist_model = MNIST_CNN().to(device)
# 设置优化器
optimizer = optim.SGD(mnist_model.parameters(), lr=MNIST_Training_Parameters['learning_rate'],
momentum=MNIST_Training_Parameters['momentum'], weight_decay=MNIST_Training_Parameters['decay'], nesterov=True)
# 训练
# 最好的验证集精度
best_val_acc = None
# 训练模型参数保存路径：/MNIST/model/MNIST_raw.pt
# **************不同模型需要修改名称**************
model_saver = '../data/MNIST/model/MART_MNIST_' + 'raw' + '.pt'
# 进行epoch次循环训练
for epoch in range(MNIST_Training_Parameters['num_epochs']):
# 一次epoch训练
train_one_epoch(model=mnist_model, train_loader=train_loader, optimizer=optimizer, epoch=epoch, device=device)
# 验证集的精度
val_acc = validation_evaluation(model=mnist_model, validation_loader=valid_loader, device=device)
adjust_MNIST_learning_rate(optimizer=optimizer, epoch=epoch)
# 每一次epoch后验证集的精度大于最好的精度时（移除模型保存路径），或者best_val_acc为None时，更新最佳精度，然后将模型参数重新写入保存路径中
if not best_val_acc or round(val_acc, 4) >= round(best_val_acc, 4):
if best_val_acc is not None:
os.remove(model_saver)
best_val_acc = val_acc
mnist_model.save(name=model_saver)
# 否则提示精度未发生提高
else:
print('Train Epoch{:>3}: validation dataset accuracy did not improve from {:.4f}\n'.format(epoch, best_val_acc))
# 测试
# 复制mnist_model
final_model = copy.deepcopy(mnist_model)
# 加载final_model
final_model.load(path=model_saver, device=device)
# 计算模型在测试集上面的精度并输出
accuracy = testing_evaluation(model=final_model, test_loader=test_loader, device=device)
# 打印模型在测试集上的精度
print('Finally, the ACCURACY of saved model [{}] on testing dataset is {:.2f}%\n'.format(final_model.model_name, accuracy * 100.0))
def search_best_radius(self,
validation_loader=None,
radius_min=0.0,
radius_max=1.0,
radius_step=0.01):
# validation_loader：验证集
# radius_min：半径最小值为0.0
# radius_max：半径最大值为1.0
# radius_step：半径迭代的步长为0.001
self.model.eval()
with torch.no_grad():
# val_acc：原始验证集的分类精度
val_acc = validation_evaluation(
model=self.model,
validation_loader=validation_loader,
device=self.device)
print(
'<--- original classification accuracy on validation dataset is {:.4f} --->'
.format(val_acc))
# total_step：总的迭代次数
total_step = int((radius_max - radius_min) / radius_step)
# 循环进行迭代
for index in range(total_step):
# tmp_radius：每一次的半径都加上一个步长
tmp_radius = radius_min + radius_step * (index + 1)
total = 0.0
correct = 0.0
for images, labels in validation_loader:
# rc_preds：返回RC预测的标签
# rc_labels：将RC预测标签转化为tensor
rc_preds = self.region_based_classification(
samples=images, radius=tmp_radius)
rc_labels = torch.from_numpy(rc_preds)
# 如果预测的标签和真实标签一致，正确个数加1
correct += (rc_labels == labels).sum().item()
# total：为总的验证集数目
total += labels.size(0)
# rc_acc：RC预测标签的准确率
rc_acc = correct / total
print(
'\tcurrent radius is {:.2f}, validation accuracy is {:.1f}/{:.1f}={:.5f}'
.format(tmp_radius, correct, total, rc_acc))
# 如果原始验证集的分类精度与RC模型的分类精度之差大于0.01时（RC的精度开始降低），则返回上一个迭代的半径
if (val_acc - rc_acc) > 1e-2:
return round(tmp_radius - radius_step, 2)
# 否则返回最大迭代半径1.0
return radius_max
def train_external_model_group(self, train_loader=None, validation_loader=None):
# train_loader：训练集
# validation_loader：验证集
# model_group：MNIST/CIFAR10的网络，各有4个
if self.Dataset == 'MNIST':
model_group = [MNIST_A(), MNIST_B(), MNIST_C(), MNIST_D()]
else:
model_group = [CIFAR10_A(), CIFAR10_B(), CIFAR10_C(), CIFAR10_D()]
model_group = [model.to(self.device) for model in model_group]
# 一个一个训练MNIST/CIFAR10的网络
for i in range(len(model_group)):
# 为MNIST/CIFAR10准备优化器，MNIST网络的优化器都为SGD，CIFAR10的最后一个网络优化器为Adam，其余的都为SGD优化器
if self.Dataset == "MNIST":
optimizer_external = optim.SGD(model_group[i].parameters(), lr=self.training_parameters['learning_rate'],
momentum=self.training_parameters['momentum'], weight_decay=self.training_parameters['decay'],
nesterov=True)
else:
# 可进行调整
if i == 3:
optimizer_external = optim.SGD(model_group[i].parameters(), lr=0.001, momentum=0.9, weight_decay=1e-6)
else:
optimizer_external = optim.Adam(model_group[i].parameters(), lr=self.training_parameters['lr'])
print('\nwe are training the {}-th static external model ......'.format(i))
# best_val_acc：验证集上的最佳分类精度
best_val_acc = None
# 对网络训练num_epochs次
for index_epoch in range(self.num_epochs):
# 一次完整的训练
train_one_epoch(model=model_group[i], train_loader=train_loader, optimizer=optimizer_external, epoch=index_epoch,
device=self.device)
# val_acc：每一次训练结束后，对验证集的分类精度
val_acc = validation_evaluation(model=model_group[i], validation_loader=validation_loader, device=self.device)
# 如果数据集为CIFAR10进行学习率调整
if self.Dataset == 'CIFAR10':
adjust_learning_rate(epoch=index_epoch, optimizer=optimizer_external)
assert os.path.exists('../DefenseEnhancedModels/{}'.format(self.defense_name))
# defense_external_saver：最佳模型参数保存位置为DefenseEnhancedModels/NEAT/MNIST_NEAT_0.pt中，或者为CIFAR10
defense_external_saver = '../DefenseEnhancedModels/{}/{}_NEAT_{}.pt'.format(self.defense_name, self.Dataset, str(i))
# 计算每次训练对验证集的分类精度，将最佳的验证精度对应的模型参数进行保存
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_external_saver)
best_val_acc = val_acc
model_group[i].save(name=defense_external_saver)
else:
print('Train Epoch {:>3}: validation dataset accuracy did not improve from {:.4f}\n'.format(index_epoch, best_val_acc))
def main(args):
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_index
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Set the random seed manually for reproducibility.
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.manual_seed(args.seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(args.seed)
np.random.seed(args.seed)
random.seed(args.seed)
# 下载CIFAR10训练集，划分为训练集和测试集（然后进行分组），保存到CIFAR10文件夹下面
train_loader, valid_loader = get_cifar10_train_validate_loader(
dir_name='../data/CIFAR10/',
batch_size=CIFAR10_Training_Parameters['batch_size'],
valid_size=0.1,
shuffle=True)
# 下载CIFAR10测试集（然后进行分组），保存到CIFAR10文件夹下面
test_loader = get_cifar10_test_loader(
dir_name='../data/CIFAR10/',
batch_size=CIFAR10_Training_Parameters['batch_size'])
# 设置模型
# **************引入的模型名称**************
resnet_model = ResNet18().to(device)
# 设置优化器
optimizer = optim.Adam(resnet_model.parameters(),
lr=CIFAR10_Training_Parameters['lr'])
# 训练
# 最好的验证集精度
best_val_acc = None
# 训练模型参数保存路径：/CIFAR10/model/CIFAR10_raw.pt
# **************不同模型需要修改名称**************
model_saver = './CIFAR10/model/ResNet18_' + 'raw' + '.pt'
# 进行epoch次循环训练
for epoch in range(CIFAR10_Training_Parameters['num_epochs']):
# 一次epoch训练
train_one_epoch(model=resnet_model,
train_loader=train_loader,
optimizer=optimizer,
epoch=epoch,
device=device)
# 验证集的精度
val_acc = validation_evaluation(model=resnet_model,
validation_loader=valid_loader,
device=device)
# 学习率调整
adjust_learning_rate(optimizer=optimizer, epoch=epoch)
# 每一次epoch后验证集的精度大于最好的精度时（移除模型保存路径），或者best_val_acc为None时，更新最佳精度，然后将模型参数重新写入保存路径中
if not best_val_acc or round(val_acc, 4) >= round(best_val_acc, 4):
if best_val_acc is not None:
os.remove(model_saver)
best_val_acc = val_acc
resnet_model.save(name=model_saver)
# 否则提示精度未发生提高
else:
print(
'Train Epoch{:>3}: validation dataset accuracy did not improve from {:.4f}\n'
.format(epoch, best_val_acc))
# 测试
# 复制resnet_model
final_model = copy.deepcopy(resnet_model)
# 加载final_model
final_model.load(path=model_saver, device=device)
# 计算模型在测试集上面的精度并输出
accuracy = testing_evaluation(model=final_model,
test_loader=test_loader,
device=device)
# 打印模型在测试集上的精度
print(
'Finally, the ACCURACY of saved model [{}] on testing dataset is {:.2f}%\n'
.format(final_model.model_name, accuracy * 100.0))
def defense(self, train_loader=None, validation_loader=None):
# train_loader：训练集
# validation_loader：验证集
# best_val_acc：验证集最佳分类精度
best_val_acc = None
# 进行num_epochs次PGD对抗训练
for epoch in range(self.num_epochs):
# 进行一次完整的RLFAT对抗训练
self.train_one_epoch_with_rbs_pgd(train_loader=train_loader,
epoch=epoch)
# val_acc：对验证集进行评估的分类精度
val_acc = validation_evaluation(
model=self.model,
validation_loader=validation_loader,
device=self.device)
assert os.path.exists('../DefenseEnhancedModels/{}'.format(
self.defense_name))
# defense_enhanced_saver：对抗训练网络参数的存放位置为DefenseEnhancedModels/PAT/CIFAR10_PAT_enhanced.pt中
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))
pgd_sense_val = 0
gamma_sense_val = 0
count = 0
# 产生白盒攻击的PGD对抗样本
for index, (images, labels) in enumerate(validation_loader):
count = index
nat_images = images.to(self.device)
nat_labels = labels.to(self.device)
self.model.eval()
# natural_images：将干净样本转化为numpy的形式
natural_images = nat_images.cpu().numpy()
# copy_images：复制干净样本numpy形式
copy_images = natural_images.copy()
# 改变图像的曝光度
gamma_copy_images = natural_images.copy()
gamma_images = batch_brightness(gamma_copy_images, c=-0.15)
gamma_images = torch.from_numpy(gamma_images).to(self.device)
# 从[-epsilon,epsilon）的均匀分布中随机采样，形成初始的扰动叠加在干净样本上
copy_images = copy_images + np.random.uniform(
-self.epsilon, self.epsilon,
copy_images.shape).astype('float32')
# 进行迭代，求PGD对抗样本
x_adv = []
for i in range(self.attack_step_num):
# 将初始的扰动样本由numpy形式转化为tensor形式
var_copy_images = torch.from_numpy(copy_images).to(
self.device)
# 可以对其进行求导
var_copy_images.requires_grad = True
# 将其输入模型进行预测
preds = self.model(var_copy_images)
# 计算损失值
loss = F.cross_entropy(preds, nat_labels)
# 对输入求梯度
gradient = torch.autograd.grad(loss, var_copy_images)[0]
# 对梯度求符号函数并转为numpy形式
gradient_sign = torch.sign(gradient).cpu().numpy()
# 对样本添加一小步扰动
copy_images = copy_images + self.step_size * gradient_sign
# 将样本的扰动大小控制在[natural_images-epsilon,natural_images+epsilon]的范围之内
copy_images = np.clip(copy_images,
natural_images - self.epsilon,
natural_images + self.epsilon)
# 将扰动大小控制在[0.0,1.0]之间
copy_images = np.clip(copy_images, 0.0, 1.0)
# 一组PGD对抗样本
pgd_adv_images = torch.from_numpy(copy_images).to(self.device)
# 计算ϵ-邻域损失灵敏度
self.model.eval()
pgd_preds = self.model(pgd_adv_images)
# 计算损失值
pgd_sense_loss = F.cross_entropy(pgd_preds, nat_labels)
# 对输入求梯度
pgd_grad = torch.autograd.grad(pgd_sense_loss,
pgd_adv_images)[0]
# 计算一个分组ϵ-邻域损失灵敏度
pgd_sense = torch.sum(pgd_grad.mul(pgd_grad))
# 计算所有分组的损失灵敏度之和
pgd_sense_val = pgd_sense_val + pgd_sense
# 计算gamma损失灵敏度
self.model.eval()
gamma_preds = self.model(gamma_images)
# 计算损失值
gamma_sense_loss = F.cross_entropy(gamma_preds, nat_labels)
# 对输入求梯度
gamma_grad = torch.autograd.grad(gamma_sense_loss,
gamma_images)[0]
# 计算一个分组的gamma灵敏度
gamma_sense = torch.sum(pgd_grad.mul(gamma_grad))
# 计算所有分组的损失灵敏度之和
gamma_sense_val = gamma_sense_val + gamma_sense
# 打印ϵ-邻域损失灵敏度和gamma损失灵敏度
if not best_val_acc or round(val_acc, 4) >= round(best_val_acc, 4):
print('ϵ loss sense val:', pgd_sense_val / (count + 1))
print('gamma loss sense val:', gamma_sense_val / (count + 1))
