def advtrain(model, device, train_loader, optimizer, epoch, log_interval):
model.train()
avg_loss = 0
# in training loop:
adversary = FGSM(model,
loss_fn=nn.NLLLoss(reduction='sum'),
eps=0.3,
clip_min=0.,
clip_max=1.,
targeted=False)
for batch_idx, (data, target) in enumerate(train_loader):
# gpu나 cpu로 데이터를 옮김
data, target = data.to(device), target.to(device)
data = adversary.perturb(data, target)
# gradient descent전에 gradient buffer를 0으로 초기화 하는 역할을 한다
optimizer.zero_grad()
output = model(data)
# negative log-likelihood: nll loss, 딥러닝 모델의 손실함수이다
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step(
) # Does the update based on the current gradient stored in grad
# 여기서 기존에 저장되어있던 gradient를 기반으로 update 하기 때문에 위의 초기화가 필요함
avg_loss += F.nll_loss(output, target, reduction='sum').item()
if batch_idx % log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
avg_loss /= len(train_loader.dataset)
return avg_loss
def create_adv_input(self, x, y, model):
# Prepare copied model
model = copy.deepcopy(model)
# Prepare input and corresponding label
data = torch.from_numpy(np.expand_dims(x, axis=0).astype(np.float32))
target = torch.from_numpy(np.array([y]).astype(np.int64))
data.requires_grad = True
from advertorch.attacks import GradientSignAttack
adversary = GradientSignAttack(model.forward, eps=self.eps)
perturbed_data = adversary.perturb(data, target)
# Have to be different
output = model.forward(perturbed_data)
final_pred = output.max(
1, keepdim=True)[1] # get the index of the max log-probability
if final_pred.item() == target.item():
return perturbed_data, 0
else:
return perturbed_data, 1
def get_attackers(attack_type: str, cfg: dict, loss_fn: nn, model: models,
*args, **kwargs):
attackers = {}
if attack_type == 'gsa' or attack_type == 'all':
# https://arxiv.org/abs/1412.652
from advertorch.attacks import GradientSignAttack
adversary = GradientSignAttack(model,
loss_fn=loss_fn,
eps=cfg['adv_gsa_eps'],
clip_min=cfg['adv_gsa_clip_min'],
clip_max=cfg['adv_gsa_clip_max'],
targeted=cfg['adv_gsa_targeted'])
attackers['gsa'] = adversary
if attack_type == 'linfpgd' or attack_type == 'all':
from advertorch.attacks import LinfPGDAttack
adversary = LinfPGDAttack(model,
loss_fn=loss_fn,
eps=cfg['adv_linfpgd_eps'],
nb_iter=cfg['adv_linfpgd_nb_iter'],
eps_iter=cfg['adv_linfpgd_eps_iter'],
rand_init=cfg['adv_linfpgd_rand_int'],
clip_min=cfg['adv_linfpgd_clip_min'],
clip_max=cfg['adv_linfpgd_clip_max'],
targeted=cfg['adv_linfpgd_targeted'])
attackers['linfpgd'] = adversary
if attack_type == 'singlepixel' or attack_type == 'all':
# https://arxiv.org/pdf/1612.06299.pdf
from advertorch.attacks import SinglePixelAttack
adversary = SinglePixelAttack(
model,
loss_fn=loss_fn,
max_pixels=cfg['adv_singlepixel_max_pixel'],
clip_min=cfg['adv_singlepixel_clip_min'],
clip_max=cfg['adv_singlepixel_clip_max'],
targeted=cfg['adv_singlepixel_targeted'])
attackers['singlepixel'] = adversary
# if attack_type=='jacobiansaliencymap' or attack_type=='all':
if attack_type == 'jacobiansaliencymap':
# https://arxiv.org/abs/1511.07528v1
from advertorch.attacks import JacobianSaliencyMapAttack
adversary = JacobianSaliencyMapAttack(
model,
num_classes=cfg['adv_jacobiansaliencymap_num_classes'],
clip_min=cfg['adv_jacobiansaliencymap_clip_min'],
clip_max=cfg['adv_jacobiansaliencymap_clip_max'],
gamma=cfg['adv_jacobiansaliencymap_gamma'],
theta=cfg['adv_jacobiansaliencymap_theta'])
attackers['jacobiansaliencymap'] = adversary
return attackers
def fgsm_attack(model, hps):
eps_list = [0., 0.1, 0.2, 0.3, 0.4, 0.5]
print('============== FGSM Summary ===============')
for eps in eps_list:
adversary = GradientSignAttack(
model,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=eps,
clip_min=-1.,
clip_max=1.,
targeted=hps.targeted)
print('epsilon = {:.4f}'.format(adversary.eps))
#attack_run(model, adversary, hps)
attack_run_rejection_policy(model, adversary, hps)
print('============== FGSM Summary ===============')
def make_adversary_dict(model, model_name, targetted=False):
if (model_name == "capsnet"):
model_for_adversary = Model_for_Adversary_Caps(model)
else:
model_for_adversary = Model_for_Adversary_CNN(model)
linf_eps = 0.3
fgsm_step = 0.05
bim_pgd_step = 0.01
adversary_dict = {}
adversary_dict['Clean'] = CleanAttack(clip_min=-0.4242, clip_max=2.8215)
adversary_dict['PGD'] = LinfPGDAttack(
model_for_adversary,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=(linf_eps / 0.3081),
nb_iter=100,
eps_iter=(bim_pgd_step / 0.3081),
rand_init=True,
clip_min=-0.4242,
clip_max=2.8215,
targeted=targetted)
adversary_dict['FGSM'] = GradientSignAttack(
model_for_adversary,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=(fgsm_step / 0.3081),
clip_min=-0.4242,
clip_max=2.8215,
targeted=targetted)
adversary_dict['BIM'] = LinfBasicIterativeAttack(
model_for_adversary,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=(linf_eps / 0.3081),
nb_iter=100,
eps_iter=(bim_pgd_step / 0.3081),
clip_min=-0.4242,
clip_max=2.8215,
targeted=targetted)
return adversary_dict
def fgsm_attack(sdim, args):
thresholds1, thresholds2 = extract_thresholds(sdim, args)
eps_list = [0.01, 0.02, 0.05, 0.1] # same as baseline DeepBayes
results_dict = {
'reject_rate1': [],
'reject_rate2': [],
'l2_distortion': []
}
for eps in eps_list:
adversary = GradientSignAttack(
sdim,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=eps,
clip_min=0.,
clip_max=1.,
targeted=args.targeted)
logger.info('epsilon = {:.4f}'.format(adversary.eps))
l2_dist, rj_rate1, rj_rate2 = adv_eval_with_rejection(
sdim, adversary, args, thresholds1, thresholds2)
results_dict['reject_rate1'].append(rj_rate1)
results_dict['reject_rate2'].append(rj_rate2)
results_dict['l2_distortion'].append(l2_dist)
torch.save(results_dict, '{}_results.pth'.format(args.attack))
batch_size = 10
AlexNet = AlexNet(3, 10).cuda()
AlexNet.load_state_dict(torch.load('./file/cifar_AlexNet.pt'))
transform = transforms.Compose([
transforms.ToTensor(),
])
train_data_set = torchvision.datasets.CIFAR10(root=cifar10Path, train=True, transform=transform, download=True)
train_loader = DataLoader(dataset=train_data_set, batch_size=batch_size, shuffle=True)
test_data_set = torchvision.datasets.CIFAR10(root=cifar10Path, train=False, transform=transform, download=True)
test_loader = DataLoader(dataset=test_data_set, batch_size=batch_size, shuffle=False)
# FGSM
adversary = GradientSignAttack(AlexNet, loss_fn=nn.CrossEntropyLoss(reduction="sum"), eps=0.004, clip_min=0., clip_max=1., targeted=True)
def generate_sample(loader, path_data, path_label):
data_adv = []
label_adv = []
i = 0
for step, (true_data, true_label) in enumerate(loader):
true_data = Variable(true_data).cuda()
true_label = Variable(true_label).cuda()
data = true_data
i += 1
if i % 2 == 1:
for j in range(20):
adv_targeted = adversary.perturb(data, (true_label + 1) % 10)
data = adv_targeted
model = resnet(num_classes=10, depth=110).cuda()
model.load_state_dict(
_preprocess_state_dict(
torch.load('./models/CIFAR10_PCL.pth.tar')['state_dict']))
model.eval()
sub = resnet(num_classes=10, depth=110).cuda()
sub.load_state_dict(
_preprocess_state_dict(
torch.load('./substitute_models/cifar_pcl_defence.pt')))
sub.eval()
adversaries = [
GradientSignAttack(model,
nn.CrossEntropyLoss(size_average=False),
eps=float(8.0 / 255)),
GradientSignAttack(sub,
nn.CrossEntropyLoss(size_average=False),
eps=float(8.0 / 255)),
]
_, _, test_loader = get_cifar10_data_loaders()
for adversary in adversaries:
correct_adv = 0
for i, (x_batch, y_batch) in enumerate(test_loader):
x_batch, y_batch = x_batch.cuda(), y_batch.cuda()
adv_x_batch = adversary.perturb(x_batch, y_batch)
logits = model(adv_x_batch)
_, preds = torch.max(logits, 1)
correct_adv += (preds == y_batch).sum().item()
progress(i + 1, len(test_loader),
if args.attack_method == "PGD":
adversary = LinfPGDAttack(
model,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=args.epsilon,
nb_iter=40,
eps_iter=0.01,
rand_init=True,
clip_min=0.0,
clip_max=1.0,
targeted=False)
elif args.attack_method == "FGSM":
adversary = GradientSignAttack(
model,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
clip_min=0.0,
clip_max=1.0,
eps=0.007,
targeted=False)
# elif args.attack_method == "JSMA":
# adversary =JacobianSaliencyMapAttack(
# model,num_classes=args.num_classes,
# clip_min=0.0, clip_max=1.0,gamma=0.145,theta=1)
elif args.attack_method == "Momentum":
adversary = MomentumIterativeAttack(
model,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=args.epsilon,
nb_iter=40,
decay_factor=1.0,
eps_iter=1.0,
if args.attack == 'PGD':
adversary = PGDAttack(lambda x: wrapper(x, pcl=pcl),
eps=epsilon,
eps_iter=epsilon / 4,
nb_iter=10,
ord=norm,
rand_init=True)
elif args.attack == 'MIFGSM':
adversary = MomentumIterativeAttack(
lambda x: wrapper(normalize(x), pcl=pcl),
eps=epsilon,
eps_iter=epsilon / 10,
ord=norm,
nb_iter=10)
elif args.attack == 'FGSM':
adversary = GradientSignAttack(lambda x: wrapper(x, pcl=pcl), eps=epsilon)
# adversary = PGDAttack(lambda x: wrapper(x, pcl=pcl), eps=epsilon, eps_iter=epsilon, nb_iter=1, ord=norm, rand_init=False)
elif args.attack == 'CW':
adversary = CarliniWagnerL2Attack(lambda x: wrapper(x, pcl=pcl),
10,
binary_search_steps=2,
max_iterations=500,
initial_const=1e-1)
elif args.attack == 'DDN':
adversary = DDN(steps=100, device=device)
ddn = True
else:
adversary = None
criterion = torch.nn.CrossEntropyLoss()
net.eval()
def _get_test_adv(attack_method,epsilon):
# define parameter
parser = argparse.ArgumentParser(description='Train MNIST')
parser.add_argument('--seed', default=0, type=int)
parser.add_argument('--mode', default="adv", help="cln | adv")
parser.add_argument('--sigma', default=75, type=int, help='noise level')
parser.add_argument('--train_batch_size', default=50, type=int)
parser.add_argument('--test_batch_size', default=1000, type=int)
parser.add_argument('--log_interval', default=200, type=int)
parser.add_argument('--result_dir', default='results', type=str, help='directory of test dataset')
parser.add_argument('--monitor', default=False, type=bool, help='if monitor the training process')
parser.add_argument('--start_save', default=90, type=int,
help='the threshold epoch which will start to save imgs data using in testing')
# attack
parser.add_argument("--attack_method", default="PGD", type=str,
choices=['FGSM', 'PGD', 'Momentum', 'STA'])
parser.add_argument('--epsilon', type=float, default=8 / 255, help='if pd_block is used')
parser.add_argument('--dataset', default='cifar10', type=str, help='dataset = [cifar10/MNIST]')
# net
parser.add_argument('--net_type', default='wide-resnet', type=str, help='model')
parser.add_argument('--depth', default=28, type=int, help='depth of model')
parser.add_argument('--widen_factor', default=10, type=int, help='width of model')
parser.add_argument('--dropout', default=0.3, type=float, help='dropout_rate')
parser.add_argument('--num_classes', default=10, type=int)
args = parser.parse_args()
torch.manual_seed(args.seed)
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
# load basic data
# 测试包装的loader
test_loader = get_handled_cifar10_test_loader(num_workers=4, shuffle=False, batch_size=50)
# 加载网络模型
# Load checkpoint
print('| Resuming from checkpoint...')
assert os.path.isdir('checkpoint'), 'Error: No checkpoint directory found!'
_, file_name = getNetwork(args)
checkpoint = torch.load('./checkpoint/' + args.dataset + os.sep + file_name + '.t7') # os.sep提供跨平台的分隔符
model = checkpoint['net']
#
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-4)
# 定义对抗攻击类型：C&W
from advertorch.attacks import LinfPGDAttack
if attack_method == "PGD":
adversary = LinfPGDAttack(
model, loss_fn=nn.CrossEntropyLoss(reduction="sum"), eps=epsilon,
nb_iter=40, eps_iter=0.01, rand_init=True, clip_min=0.0, clip_max=1.0,
targeted=False)
elif attack_method == "FGSM":
adversary = GradientSignAttack(
model, loss_fn=nn.CrossEntropyLoss(reduction="sum"),
clip_min=0.0, clip_max=1.0, eps=0.007, targeted=False) # 先测试一下不含扰动范围限制的，FGSM的eps代表的是一般的eps_iter
elif attack_method == "Momentum":
adversary = MomentumIterativeAttack(
model, loss_fn=nn.CrossEntropyLoss(reduction="sum"), eps=epsilon,
nb_iter=40, decay_factor=1.0, eps_iter=1.0, clip_min=0.0, clip_max=1.0,
targeted=False, ord=np.inf)
elif attack_method == "STA":
adversary = SpatialTransformAttack(
model, num_classes=args.num_classes, loss_fn=nn.CrossEntropyLoss(reduction="sum"),
initial_const=0.05, max_iterations=1000, search_steps=1, confidence=0, clip_min=0.0, clip_max=1.0,
targeted=False, abort_early=True) # 先测试一下不含扰动范围限制的
# generate for train.h5 | save as train_adv_attackMethod_epsilon
test_adv = []
test_true_target = []
for clndata, target in test_loader:
print("clndata:{}".format(clndata.size()))
clndata, target = clndata.to(device), target.to(device)
with ctx_noparamgrad_and_eval(model):
advdata = adversary.perturb(clndata, target)
test_adv.append(advdata.detach().cpu().numpy())
test_true_target.append(target.cpu().numpy())
test_adv = np.reshape(np.asarray(test_adv),[-1,3,32,32])
test_true_target = np.reshape(np.asarray(test_true_target),[-1])
print("test_adv.shape:{}".format(test_adv.shape))
print("test_true_target.shape:{}".format(test_true_target.shape))
del model
return test_adv, test_true_target
def _get_test_adv(attack_method, epsilon):
torch.manual_seed(args.seed)
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
# load basic data
# 测试包装的loader
test_loader = get_handled_cifar10_test_loader(num_workers=4,
shuffle=False,
batch_size=50)
# 加载网络模型
# Load checkpoint
print('| Resuming from checkpoint...')
assert os.path.isdir('checkpoint'), 'Error: No checkpoint directory found!'
_, file_name = getNetwork(args)
checkpoint = torch.load('./checkpoint/' + args.dataset + os.sep +
file_name + '.t7') # os.sep提供跨平台的分隔符
model = checkpoint['net']
#
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-4)
# 定义对抗攻击类型：C&W
from advertorch.attacks import LinfPGDAttack
if attack_method == "PGD":
adversary = LinfPGDAttack(model,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=epsilon,
nb_iter=20,
eps_iter=0.01,
rand_init=True,
clip_min=0.0,
clip_max=1.0,
targeted=False)
elif attack_method == "FGSM":
adversary = GradientSignAttack(
model,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
clip_min=0.0,
clip_max=1.0,
eps=epsilon,
targeted=False) # 先测试一下不含扰动范围限制的，FGSM的eps代表的是一般的eps_iter
elif attack_method == "Momentum":
adversary = MomentumIterativeAttack(
model,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=epsilon,
nb_iter=20,
decay_factor=1.0,
eps_iter=1.0,
clip_min=0.0,
clip_max=1.0,
targeted=False,
ord=np.inf)
elif attack_method == "STA":
adversary = SpatialTransformAttack(
model,
num_classes=args.num_classes,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
initial_const=0.05,
max_iterations=1000,
search_steps=1,
confidence=0,
clip_min=0.0,
clip_max=1.0,
targeted=False,
abort_early=True) # 先测试一下不含扰动范围限制的
# generate for train.h5 | save as train_adv_attackMethod_epsilon
test_adv = []
test_true_target = []
for clndata, target in test_loader:
print("clndata:{}".format(clndata.size()))
clndata, target = clndata.to(device), target.to(device)
with ctx_noparamgrad_and_eval(model):
advdata = adversary.perturb(clndata, target)
test_adv.append(advdata.detach().cpu().numpy())
test_true_target.append(target.cpu().numpy())
test_adv = np.reshape(np.asarray(test_adv), [-1, 3, 32, 32])
test_true_target = np.reshape(np.asarray(test_true_target), [-1])
print("test_adv.shape:{}".format(test_adv.shape))
print("test_true_target.shape:{}".format(test_true_target.shape))
del model
return test_adv, test_true_target
def main():
# get args
args = get_args()
# set up gpus
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
assert torch.cuda.is_available()
# load models
if 'gal' in args.model_file:
leaky_relu = True
else:
leaky_relu = False
ensemble = utils.get_models(args, train=False, as_ensemble=True, model_file=args.model_file, leaky_relu=leaky_relu)
# get data loaders
total_sample_num = 10000
if args.subset_num:
random.seed(0)
subset_idx = random.sample(range(total_sample_num), args.subset_num)
testloader = utils.get_testloader(args, batch_size=200, shuffle=False, subset_idx=subset_idx)
else:
testloader = utils.get_testloader(args, batch_size=200, shuffle=False)
loss_fn = nn.CrossEntropyLoss() if args.loss_fn == 'xent' else CarliniWagnerLoss(conf=args.cw_conf)
rob = {}
rob['sample_num'] = args.subset_num if args.subset_num else total_sample_num
rob['loss_fn'] = 'xent' if args.loss_fn == 'xent' else 'cw_{:.1f}'.format(args.cw_conf)
train_seed = args.model_file.split('/')[-3]
train_alg = args.model_file.split('/')[-4]
if args.convergence_check:
eps = 0.01
steps_list = [50, 500, 1000]
random_start = 1
rob['random_start'] = random_start
rob['eps'] = eps
# FGSM
test_iter = tqdm(testloader, desc='FGSM', leave=False, position=0)
adversary = GradientSignAttack(
ensemble, loss_fn=nn.CrossEntropyLoss(), eps=eps,
clip_min=0., clip_max=1., targeted=False)
_, label, pred, advpred = attack_whole_dataset(adversary, test_iter, device="cuda")
print("Accuracy: {:.2f}%, FGSM Accuracy: {:.2f}%".format(
100. * (label == pred).sum().item() / len(label),
100. * (label == advpred).sum().item() / len(label)))
rob['clean'] = 100. * (label == pred).sum().item() / len(label)
rob['fgsm'] = 100. * (label == advpred).sum().item() / len(label)
for steps in tqdm(steps_list, desc='PGD steps', leave=False, position=0):
correct_or_not = []
for i in tqdm(range(random_start), desc='Random Start', leave=False, position=1):
torch.manual_seed(i)
test_iter = tqdm(testloader, desc='Batch', leave=False, position=2)
adversary = LinfPGDAttack(
ensemble, loss_fn=loss_fn, eps=eps,
nb_iter=steps, eps_iter=eps/5, rand_init=True, clip_min=0., clip_max=1.,
targeted=False)
_, label, pred, advpred = attack_whole_dataset(adversary, test_iter, device="cuda")
correct_or_not.append(label == advpred)
correct_or_not = torch.stack(correct_or_not, dim=-1).all(dim=-1)
tqdm.write("Accuracy: {:.2f}%, steps: {:d}, PGD Accuracy: {:.2f}%".format(
100. * (label == pred).sum().item() / len(label),
steps,
100. * correct_or_not.sum().item() / len(label)))
rob[str(steps)] = 100. * correct_or_not.sum().item() / len(label)
# save to file
if args.save_to_csv:
output_root = os.path.join('results', 'wbox', train_alg, train_seed, 'convergence_check')
if not os.path.exists(output_root):
os.makedirs(output_root)
output_filename = args.model_file.split('/')[-2]
output = os.path.join(output_root, '.'.join((output_filename, 'csv')))
df = pd.DataFrame(rob, index=[0])
if args.append_out and os.path.isfile(output):
with open(output, 'a') as f:
f.write('\n')
df.to_csv(output, sep=',', mode='a', header=False, index=False, float_format='%.2f')
else:
df.to_csv(output, sep=',', index=False, float_format='%.2f')
else:
eps_list = [0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07]
rob['random_start'] = args.random_start
rob['steps'] = args.steps
for eps in tqdm(eps_list, desc='PGD eps', leave=True, position=0):
correct_or_not = []
for i in tqdm(range(args.random_start), desc='Random Start', leave=False, position=1):
torch.manual_seed(i)
test_iter = tqdm(testloader, desc='Batch', leave=False, position=2)
adversary = LinfPGDAttack(
ensemble, loss_fn=loss_fn, eps=eps,
nb_iter=args.steps, eps_iter=eps/5, rand_init=True, clip_min=0., clip_max=1.,
targeted=False)
_, label, pred, advpred = attack_whole_dataset(adversary, test_iter, device="cuda")
correct_or_not.append(label == advpred)
correct_or_not = torch.stack(correct_or_not, dim=-1).all(dim=-1)
tqdm.write("Accuracy: {:.2f}%, eps: {:.2f}, PGD Accuracy: {:.2f}%".format(
100. * (label == pred).sum().item() / len(label),
eps,
100. * correct_or_not.sum().item() / len(label)))
rob['clean'] = 100. * (label == pred).sum().item() / len(label)
rob[str(eps)] = 100. * correct_or_not.sum().item() / len(label)
# save to file
if args.save_to_csv:
output_root = os.path.join('results', 'wbox', train_alg, train_seed)
if not os.path.exists(output_root):
os.makedirs(output_root)
output_filename = args.model_file.split('/')[-2]
output = os.path.join(output_root, '.'.join((output_filename, 'csv')))
df = pd.DataFrame(rob, index=[0])
if args.append_out and os.path.isfile(output):
with open(output, 'a') as f:
f.write('\n')
df.to_csv(output, sep=',', mode='a', header=False, index=False, float_format='%.2f')
else:
df.to_csv(output, sep=',', index=False, float_format='%.2f')
def train_Ours(args, train_loader, val_loader, knownclass, Encoder, Decoder,
NorClsfier, SSDClsfier, summary_writer, saver):
seed = init_random_seed(args.manual_seed)
criterionCls = nn.CrossEntropyLoss()
criterionRec = nn.MSELoss()
if args.parallel_train:
Encoder = DataParallel(Encoder)
Decoder = DataParallel(Decoder)
NorClsfier = DataParallel(NorClsfier)
SSDClsfier = DataParallel(SSDClsfier)
optimizer = optim.Adam(
list(Encoder.parameters()) + list(NorClsfier.parameters()) +
list(SSDClsfier.parameters()) + list(Decoder.parameters()),
lr=args.lr)
if args.adv is 'PGDattack':
print("**********Defense PGD Attack**********")
elif args.adv is 'FGSMattack':
print("**********Defense FGSM Attack**********")
if args.adv is 'PGDattack':
from advertorch.attacks import PGDAttack
nor_adversary = PGDAttack(predict1=Encoder,
predict2=NorClsfier,
nb_iter=args.adv_iter)
rot_adversary = PGDAttack(predict1=Encoder,
predict2=SSDClsfier,
nb_iter=args.adv_iter)
elif args.adv is 'FGSMattack':
from advertorch.attacks import GradientSignAttack
nor_adversary = GradientSignAttack(predict1=Encoder,
predict2=NorClsfier)
rot_adversary = GradientSignAttack(predict1=Encoder,
predict2=SSDClsfier)
global_step = 0
# ----------
# Training
# ----------
for epoch in range(args.n_epoch):
Encoder.train()
Decoder.train()
NorClsfier.train()
SSDClsfier.train()
for steps, (orig, label, rot_orig,
rot_label) in enumerate(train_loader):
label = lab_conv(knownclass, label)
orig, label = orig.cuda(), label.long().cuda()
rot_orig, rot_label = rot_orig.cuda(), rot_label.long().cuda()
with ctx_noparamgrad_and_eval(Encoder):
with ctx_noparamgrad_and_eval(NorClsfier):
with ctx_noparamgrad_and_eval(SSDClsfier):
adv = nor_adversary.perturb(orig, label)
rot_adv = rot_adversary.perturb(rot_orig, rot_label)
latent_feat = Encoder(adv)
norpred = NorClsfier(latent_feat)
norlossCls = criterionCls(norpred, label)
recon = Decoder(latent_feat)
lossRec = criterionRec(recon, orig)
ssdpred = SSDClsfier(Encoder(rot_adv))
rotlossCls = criterionCls(ssdpred, rot_label)
loss = args.norClsWgt * norlossCls + args.rotClsWgt * rotlossCls + args.RecWgt * lossRec
optimizer.zero_grad()
loss.backward()
optimizer.step()
#============ tensorboard the log info ============#
lossinfo = {
'loss': loss.item(),
'norlossCls': norlossCls.item(),
'lossRec': lossRec.item(),
'rotlossCls': rotlossCls.item(),
}
global_step += 1
#============ print the log info ============#
if (steps + 1) % args.log_step == 0:
errors = OrderedDict([
('loss', loss.item()),
('norlossCls', norlossCls.item()),
('lossRec', lossRec.item()),
('rotlossCls', rotlossCls.item()),
])
saver.print_current_errors((epoch + 1), (steps + 1), errors)
# evaluate performance on validation set periodically
if ((epoch + 1) % args.val_epoch == 0):
# switch model to evaluation mode
Encoder.eval()
NorClsfier.eval()
running_corrects = 0.0
epoch_size = 0.0
val_loss_list = []
# calculate accuracy on validation set
for steps, (images, label) in enumerate(val_loader):
label = lab_conv(knownclass, label)
images, label = images.cuda(), label.long().cuda()
adv = nor_adversary.perturb(images, label)
with torch.no_grad():
logits = NorClsfier(Encoder(adv))
_, preds = torch.max(logits, 1)
running_corrects += torch.sum(preds == label.data)
epoch_size += images.size(0)
val_loss = criterionCls(logits, label)
val_loss_list.append(val_loss.item())
val_loss_mean = sum(val_loss_list) / len(val_loss_list)
val_acc = running_corrects.double() / epoch_size
print('Val Acc: {:.4f}, Val Loss: {:.4f}'.format(
val_acc, val_loss_mean))
valinfo = {
'Val Acc': val_acc.item(),
'Val Loss': val_loss.item(),
}
for tag, value in valinfo.items():
summary_writer.add_scalar(tag, value, (epoch + 1))
orig_show = vutils.make_grid(orig, normalize=True, scale_each=True)
recon_show = vutils.make_grid(recon,
normalize=True,
scale_each=True)
summary_writer.add_image('Ori_Image', orig_show, (epoch + 1))
summary_writer.add_image('Rec_Image', recon_show, (epoch + 1))
if ((epoch + 1) % args.model_save_epoch == 0):
model_save_path = os.path.join(args.results_path,
args.training_type, 'snapshots',
args.datasetname + '-' + args.split,
args.denoisemean,
args.adv + str(args.adv_iter))
mkdir(model_save_path)
torch.save(
Encoder.state_dict(),
os.path.join(model_save_path,
"Encoder-{}.pt".format(epoch + 1)))
torch.save(
NorClsfier.state_dict(),
os.path.join(model_save_path,
"NorClsfier-{}.pt".format(epoch + 1)))
torch.save(
Decoder.state_dict(),
os.path.join(model_save_path,
"Decoder-{}.pt".format(epoch + 1)))
torch.save(Encoder.state_dict(),
os.path.join(model_save_path, "Encoder-final.pt"))
torch.save(NorClsfier.state_dict(),
os.path.join(model_save_path, "NorClsfier-final.pt"))
torch.save(Decoder.state_dict(),
os.path.join(model_save_path, "Decoder-final.pt"))
def test_adver(net, tar_net, attack, target):
net.eval()
tar_net.eval()
# BIM
if attack == 'BIM':
adversary = LinfBasicIterativeAttack(
net,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=0.25,
nb_iter=120,
eps_iter=0.02,
clip_min=0.0,
clip_max=1.0,
targeted=opt.target)
# PGD
elif attack == 'PGD':
if opt.target:
adversary = PGDAttack(net,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=0.25,
nb_iter=11,
eps_iter=0.03,
clip_min=0.0,
clip_max=1.0,
targeted=opt.target)
else:
adversary = PGDAttack(net,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=0.25,
nb_iter=6,
eps_iter=0.03,
clip_min=0.0,
clip_max=1.0,
targeted=opt.target)
# FGSM
elif attack == 'FGSM':
adversary = GradientSignAttack(
net,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=0.26,
targeted=opt.target)
elif attack == 'CW':
adversary = CarliniWagnerL2Attack(
net,
num_classes=10,
learning_rate=0.45,
# loss_fn=nn.CrossEntropyLoss(reduction="sum"),
binary_search_steps=10,
max_iterations=12,
targeted=opt.target)
# ----------------------------------
# Obtain the accuracy of the model
# ----------------------------------
with torch.no_grad():
correct_netD = 0.0
total = 0.0
net.eval()
for data in testloader:
inputs, labels = data
inputs = inputs.cuda()
labels = labels.cuda()
outputs = net(inputs)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct_netD += (predicted == labels).sum()
print('Accuracy of the network on netD: %.2f %%' %
(100. * correct_netD.float() / total))
# ----------------------------------
# Obtain the attack success rate of the model
# ----------------------------------
correct = 0.0
total = 0.0
tar_net.eval()
total_L2_distance = 0.0
for data in testloader:
inputs, labels = data
inputs = inputs.to(device)
labels = labels.to(device)
outputs = tar_net(inputs)
_, predicted = torch.max(outputs.data, 1)
if target:
# randomly choose the specific label of targeted attack
labels = torch.randint(0, 9, (1, )).to(device)
# test the images which are not classified as the specific label
if predicted != labels:
# print(total)
adv_inputs_ori = adversary.perturb(inputs, labels)
L2_distance = (torch.norm(adv_inputs_ori - inputs)).item()
total_L2_distance += L2_distance
with torch.no_grad():
outputs = tar_net(adv_inputs_ori)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum()
else:
# test the images which are classified correctly
if predicted == labels:
# print(total)
adv_inputs_ori = adversary.perturb(inputs, labels)
L2_distance = (torch.norm(adv_inputs_ori - inputs)).item()
total_L2_distance += L2_distance
with torch.no_grad():
outputs = tar_net(adv_inputs_ori)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum()
if target:
print('Attack success rate: %.2f %%' %
(100. * correct.float() / total))
else:
print('Attack success rate: %.2f %%' %
(100.0 - 100. * correct.float() / total))
print('l2 distance: %.4f ' % (total_L2_distance / total))
from torch.autograd import Variable
from small_cnn import SmallCNN
from advertorch.attacks import GradientSignAttack, LinfBasicIterativeAttack
from utils import progress, get_mnist_data_loaders
import numpy as np
model = SmallCNN().cuda()
model.load_state_dict(torch.load('./models/trades.pt'))
model.eval()
sub = SmallCNN().cuda()
sub.load_state_dict(torch.load('./substitute_models/mnist_trades.pt'))
sub.eval()
adversaries = [
GradientSignAttack(model, nn.CrossEntropyLoss(size_average=False),
eps=0.3),
GradientSignAttack(sub, nn.CrossEntropyLoss(size_average=False), eps=0.3),
LinfBasicIterativeAttack(model,
nn.CrossEntropyLoss(size_average=False),
eps=0.3,
nb_iter=40,
eps_iter=0.01),
LinfBasicIterativeAttack(sub,
nn.CrossEntropyLoss(size_average=False),
eps=0.3,
nb_iter=40,
eps_iter=0.01)
]
_, _, test_loader = get_mnist_data_loaders()
for adversary in adversaries:
correct_adv = 0
def adv_train_loop(model,
params,
ds,
min_y,
base_data,
model_id,
attack_type,
device,
batch_size,
max_epochs=5):
print('training adversarial:', attack_type)
ds_train, ds_valid = ds
min_y_train, min_y_val = min_y
original_model = copy.deepcopy(
model) # used to generate adv images for the trained model
original_model.eval()
model = copy.deepcopy(
model) # making a copy so that original model is not changed
model = model.to(device)
model_id = f'{model_id}_{attack_type}'
with create_summary_writer(model,
ds_train,
base_data,
model_id,
device=device) as writer:
lr = params['lr']
mom = params['momentum']
wd = params['l2_wd']
optimizer = torch.optim.SGD(model.parameters(),
lr=lr,
momentum=mom,
weight_decay=wd)
sched = ReduceLROnPlateau(optimizer, factor=0.5, patience=5)
funcs = {'accuracy': Accuracy(), 'loss': Loss(F.cross_entropy)}
loss = funcs['loss']._loss_fn
acc_metric = Accuracy(device=device)
loss_metric = Loss(F.cross_entropy, device=device)
acc_val_metric = Accuracy(device=device)
loss_val_metric = Loss(F.cross_entropy, device=device)
classifier = PyTorchClassifier(
model=original_model,
clip_values=(0, 1),
loss=nn.CrossEntropyLoss(),
optimizer=optimizer,
input_shape=(3, 64, 64),
nb_classes=200,
)
attack = None
# if attack_type == "fgsm":
# attack = FastGradientMethod(estimator=classifier, eps=0.2)
# elif attack_type == "bim":
# attack = BasicIterativeMethod(estimator=classifier, eps=0.2)
# elif attack_type == "carlini":
# attack = CarliniLInfMethod(classifier=classifier)
# elif attack_type == "deepfool":
# attack = DeepFool(classifier=classifier)
if attack_type == "fgsm":
attack = GradientSignAttack(model, loss_fn=loss, eps=0.2)
elif attack_type == "ffa":
attack = FastFeatureAttack(model, loss_fn=loss, eps=0.3)
elif attack_type == "carlini":
attack = CarliniWagnerL2Attack(model, 200, max_iterations=1000)
elif attack_type == "lbfgs":
attack = DeepFool(classifier=classifier)
def train_step(engine, batch):
model.train()
x, y = batch
x = x.to(device)
y = y.to(device) - min_y_train
with ctx_noparamgrad_and_eval(model):
x_adv = attack.perturb(x, y)
optimizer.zero_grad()
x = torch.cat((x, x_adv))
y = torch.cat((y, y))
ans = model.forward(x)
l = loss(ans, y)
optimizer.zero_grad()
l.backward()
optimizer.step()
# return ans, y
return l.item()
trainer = Engine(train_step)
# acc_metric.attach(trainer, "accuracy")
# loss_metric.attach(trainer, 'loss')
def train_eval_step(engine, batch):
model.eval()
x, y = batch
x = x.to(device)
y = y.to(device) - min_y_train
x_adv = attack.perturb(x, y)
x = torch.cat((x, x_adv))
y = torch.cat((y, y))
with torch.no_grad():
ans = model.forward(x)
return ans, y
train_evaluator = Engine(train_eval_step)
acc_metric.attach(train_evaluator, "accuracy")
loss_metric.attach(train_evaluator, 'loss')
def validation_step(engine, batch):
model.eval()
x, y = batch
x = x.to(device)
y = y.to(device) - min_y_val
x_adv = attack.perturb(x, y)
x = torch.cat((x, x_adv))
y = torch.cat((y, y))
with torch.no_grad():
ans = model.forward(x)
return ans, y
valid_evaluator = Engine(validation_step)
acc_val_metric.attach(valid_evaluator, "accuracy")
loss_val_metric.attach(valid_evaluator, 'loss')
@trainer.on(
Events.ITERATION_COMPLETED(every=200 * 5000 // batch_size // 10))
def log_validation_results(engine):
valid_evaluator.run(ds_valid)
metrics = valid_evaluator.state.metrics
valid_avg_accuracy = metrics['accuracy']
avg_nll = metrics['loss']
print(
"Validation Results - Epoch: {} Avg accuracy: {:.2f} Avg loss: {:.2f}"
.format(engine.state.epoch, valid_avg_accuracy, avg_nll))
writer.add_scalar("validation/avg_loss", avg_nll,
engine.state.epoch)
writer.add_scalar("validation/avg_accuracy", valid_avg_accuracy,
engine.state.epoch)
writer.add_scalar("validation/avg_error", 1. - valid_avg_accuracy,
engine.state.epoch)
@trainer.on(Events.EPOCH_COMPLETED)
def lr_scheduler(engine):
metrics = valid_evaluator.state.metrics
avg_nll = metrics['accuracy']
sched.step(avg_nll)
@trainer.on(Events.ITERATION_COMPLETED(every=50))
def log_training_loss(engine):
batch = engine.state.batch
ds = DataLoader(TensorDataset(*batch), batch_size=batch_size)
train_evaluator.run(ds)
metrics = train_evaluator.state.metrics
# metrics = engine.state.metrics
accuracy = metrics['accuracy']
nll = metrics['loss']
iter = (engine.state.iteration - 1) % len(ds_train) + 1
if (iter % 50) == 0:
print("Epoch[{}] Iter[{}/{}] Accuracy: {:.2f} Loss: {:.2f}".
format(engine.state.epoch, iter, len(ds_train), accuracy,
nll))
writer.add_scalar("batchtraining/detloss", nll, engine.state.epoch)
writer.add_scalar("batchtraining/accuracy", accuracy,
engine.state.iteration)
writer.add_scalar("batchtraining/error", 1. - accuracy,
engine.state.iteration)
writer.add_scalar("batchtraining/loss", engine.state.output,
engine.state.iteration)
@trainer.on(Events.EPOCH_COMPLETED)
def log_lr(engine):
writer.add_scalar("lr", optimizer.param_groups[0]['lr'],
engine.state.epoch)
# @trainer.on(Events.EPOCH_COMPLETED)
# def log_training_results(engine):
# train_evaluator.run(ds_train)
# metrics = train_evaluator.state.metrics
# # metrics = engine.state.metrics
# avg_accuracy = metrics['accuracy']
# avg_nll = metrics['loss']
# print("Training Results - Epoch: {} Avg accuracy: {:.2f} Avg loss: {:.2f}"
# .format(engine.state.epoch, avg_accuracy, avg_nll))
# writer.add_scalar("training/avg_loss", avg_nll, engine.state.epoch)
# writer.add_scalar("training/avg_accuracy",
# avg_accuracy, engine.state.epoch)
# writer.add_scalar("training/avg_error", 1. -
# avg_accuracy, engine.state.epoch)
@trainer.on(
Events.ITERATION_COMPLETED(every=200 * 5000 // batch_size // 10))
def validation_value(engine):
metrics = valid_evaluator.state.metrics
valid_avg_accuracy = metrics['accuracy']
return valid_avg_accuracy
to_save = {'model': model}
handler = Checkpoint(
to_save,
DiskSaver(os.path.join(base_data, model_id), create_dir=True),
score_function=validation_value,
score_name="val_acc",
global_step_transform=global_step_from_engine(trainer),
n_saved=None)
# kick everything off
trainer.add_event_handler(
Events.ITERATION_COMPLETED(every=200 * 5000 // batch_size // 10),
handler)
trainer.run(ds_train, max_epochs=max_epochs)
def _generate_adv_file(attack_method, num_classes, epsilon, set_size):
# load model
model = torch.load(os.path.join("checkpoint", "resnet50_epoch_22.pth"))
model = model.cuda()
#define attack
if attack_method == "PGD":
adversary = LinfPGDAttack(model,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=epsilon,
nb_iter=20,
eps_iter=0.01,
rand_init=True,
clip_min=0.0,
clip_max=1.0,
targeted=False)
elif attack_method == "FGSM":
adversary = GradientSignAttack(
model,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
clip_min=0.0,
clip_max=1.0,
eps=epsilon,
targeted=False)
elif attack_method == "Momentum":
adversary = MomentumIterativeAttack(
model,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=epsilon,
nb_iter=20,
decay_factor=1.0,
eps_iter=1.0,
clip_min=0.0,
clip_max=1.0,
targeted=False,
ord=np.inf)
elif attack_method == "STA":
adversary = SpatialTransformAttack(
model,
num_classes=num_classes,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
initial_const=0.05,
max_iterations=500,
search_steps=1,
confidence=0,
clip_min=0.0,
clip_max=1.0,
targeted=False,
abort_early=True)
elif attack_method == "DeepFool":
adversary = DeepFool(model,
max_iter=20,
clip_max=1.0,
clip_min=0.0,
epsilon=epsilon)
elif attack_method == "CW":
adversary = CarliniWagnerL2Attack(
model,
num_classes=args.num_classes,
epsilon=epsilon,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
max_iterations=20,
confidence=0,
clip_min=0.0,
clip_max=1.0,
targeted=False,
abort_early=True)
# version two
h5_store = h5py.File("data/test_tiny_ImageNet_" + str(set_size) + ".h5",
"r")
data = h5_store['data'][:]
target = h5_store['true_target'][:]
data = torch.from_numpy(data)
target = torch.from_numpy(target)
test_dataset = ImageNetDataset(data, target)
test_loader = DataLoader(dataset=test_dataset,
num_workers=4,
drop_last=True,
batch_size=50,
shuffle=False)
torch.manual_seed(0)
test_adv = np.zeros([set_size, 3, 64, 64])
test_true_target = np.zeros([set_size])
# perturb
for batch_idx, (clndata, target) in enumerate(test_loader):
print("{}/{}".format(batch_idx, set_size // 50))
clndata, target = clndata.cuda().float(), target.cuda().long()
with ctx_noparamgrad_and_eval(model):
# print(target)
advdata = adversary.perturb(clndata, target)
test_adv[batch_idx * 50:(batch_idx + 1) *
50, :, :, :] = advdata.detach().cpu().numpy()
test_true_target[batch_idx * 50:(batch_idx + 1) *
50] = target.cpu().numpy()
print("test_adv.shape:{}".format(test_adv.shape))
print("test_true_target.shape:{}".format(test_true_target.shape))
del model
h5_store = h5py.File(
"data/test_tiny_ImageNet_" + str(set_size) + "_adv_" +
str(attack_method) + "_" + str(epsilon) + ".h5", 'w')
h5_store.create_dataset('data', data=test_adv)
h5_store.create_dataset('true_target', data=test_true_target)
h5_store.close()
model.load_state_dict(
torch.load("checkpoint/vgg_baseline.pth"))
model.to(device)
model.eval()
cifar100_test_loader = get_test_dataloader(
settings.CIFAR100_TRAIN_MEAN,
settings.CIFAR100_TRAIN_STD,
#settings.CIFAR100_PATH,
num_workers=2,
batch_size=16,
shuffle=True
)
adversary = GradientSignAttack(
model, loss_fn=nn.CrossEntropyLoss(reduction="sum"), eps=0.3,
clip_min=0.0, clip_max=1.0,
targeted=False)
correct_1 = 0.0
correct_5 = 0.0
attack_correct_1 = 0.0
attack_correct_5 = 0.0
total = 0
for n_iter, (image, label) in enumerate(cifar100_test_loader):
print("iteration: {}\ttotal {} iterations".format(n_iter + 1, len(cifar100_test_loader)))
image = Variable(image).cuda()
label = Variable(label).cuda()
output = model(image)
adv_untargeted = adversary.perturb(image,None)
attack_output = model(adv_untargeted)
# for i in range(16):
# save_image_tensor2pillow(image[i],str(i)+'.jpg')
elif mode == 'adv':
nb_epoch = 50
model_filename = "mnist_lenet5_advtrained.pth"
else:
raise RuntimeError('mode must be "cls" or "adv"')
train_loader = get_mnist_train_loader(batch_size=train_bs, shuffle=True)
test_loader = get_mnist_test_loader(batch_size=test_bs, shuffle=False)
model = LeNet5()
model.to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-4)
adversary = GradientSignAttack(model,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=0.3,
clip_min=0.0,
clip_max=1.0,
targeted=False)
for epoch in range(nb_epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
ori = data
if mode == 'adv' or (mode == 'half' and random.random() > 0.5):
# when performing attack, the model needs to be in eval mode
# also the parameters should be accumulating gradients
with ctx_noparamgrad_and_eval(model):
data = adversary.perturb(data, target)
optimizer.zero_grad()
U_sampler = torch.utils.data.SubsetRandomSampler(U_idx)
S_loader1 = torch.utils.data.DataLoader(trainset,
batch_size=S_batch_size,
sampler=S_sampler)
S_loader2 = torch.utils.data.DataLoader(trainset,
batch_size=S_batch_size,
sampler=S_sampler)
U_loader = torch.utils.data.DataLoader(trainset,
batch_size=U_batch_size,
sampler=U_sampler)
#net1 adversary object
adversary1 = GradientSignAttack(net1,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=args.epsilon,
clip_min=-math.inf,
clip_max=math.inf,
targeted=False)
#net2 adversary object
adversary2 = GradientSignAttack(net2,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=args.epsilon,
clip_min=-math.inf,
clip_max=math.inf,
targeted=False)
if args.dataset == 'cifar10':
step = int(len(trainset) / batch_size)
# for SVHN, not implemented yet
else:
def train(epoch):
net1.train()
net2.train()
adjust_learning_rate(optimizer, epoch)
adjust_lamda(epoch)
total_S1 = 0
total_S2 = 0
total_U1 = 0
total_U2 = 0
train_correct_S1 = 0
train_correct_S2 = 0
train_correct_U1 = 0
train_correct_U2 = 0
running_loss = 0.0
ls = 0.0
lc = 0.0
ld = 0.0
# create iterator for b1, b2, bu
S_iter1 = iter(S_loader1)
S_iter2 = iter(S_loader2)
U_iter = iter(U_loader)
print('epoch:', epoch + 1)
# for i, (inputs, targets) in enumerate(zip(S_loader1, S_loader2, U_loader):
for i in tqdm(range(step)):
inputs_S1, labels_S1 = S_iter1.next()
inputs_S2, labels_S2 = S_iter2.next()
inputs_U, labels_U = U_iter.next(
) # note that labels_U will not be used for training.
inputs_S1 = inputs_S1.cuda()
labels_S1 = labels_S1.cuda()
inputs_S2 = inputs_S2.cuda()
labels_S2 = labels_S2.cuda()
inputs_U = inputs_U.cuda()
logit_S1 = net1(inputs_S1)
logit_S2 = net2(inputs_S2)
logit_U1 = net1(inputs_U)
logit_U2 = net2(inputs_U)
_, predictions_S1 = torch.max(logit_S1, 1)
_, predictions_S2 = torch.max(logit_S2, 1)
# pseudo labels of U
_, predictions_U1 = torch.max(logit_U1, 1)
_, predictions_U2 = torch.max(logit_U2, 1)
#net1 adversary object
adversary1 = GradientSignAttack(
net1,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=args.epsilon,
clip_min=0.0,
clip_max=1.0,
targeted=False)
#net2 adversary object
adversary2 = GradientSignAttack(
net2,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=args.epsilon,
clip_min=0.0,
clip_max=1.0,
targeted=False)
#generate adversarial examples
perturbed_data_S1 = adversary1.perturb(inputs_S1, labels_S1)
perturbed_data_U1 = adversary1.perturb(inputs_U, predictions_U1)
perturbed_data_S2 = adversary2.perturb(inputs_S2, labels_S2)
perturbed_data_U2 = adversary2.perturb(inputs_U, predictions_U2)
perturbed_logit_S1 = net1(perturbed_data_S2)
perturbed_logit_S2 = net2(perturbed_data_S1)
perturbed_logit_U1 = net1(perturbed_data_U2)
perturbed_logit_U2 = net2(perturbed_data_U1)
# zero the parameter gradients
optimizer.zero_grad()
net1.zero_grad()
net2.zero_grad()
Loss_sup = loss_sup(logit_S1, logit_S2, labels_S1, labels_S2)
Loss_cot = loss_cot(logit_U1, logit_U2)
Loss_diff = loss_diff(logit_S1, logit_S2, perturbed_logit_S1,
perturbed_logit_S2, logit_U1, logit_U2,
perturbed_logit_U1, perturbed_logit_U2)
total_loss = Loss_sup + lambda_cot * Loss_cot + lambda_diff * Loss_diff
total_loss.backward()
optimizer.step()
train_correct_S1 += np.sum(
predictions_S1.cpu().numpy() == labels_S1.cpu().numpy())
total_S1 += labels_S1.size(0)
train_correct_U1 += np.sum(
predictions_U1.cpu().numpy() == labels_U.cpu().numpy())
total_U1 += labels_U.size(0)
train_correct_S2 += np.sum(
predictions_S2.cpu().numpy() == labels_S2.cpu().numpy())
total_S2 += labels_S2.size(0)
train_correct_U2 += np.sum(
predictions_U2.cpu().numpy() == labels_U.cpu().numpy())
total_U2 += labels_U.size(0)
running_loss += total_loss.item()
ls += Loss_sup.item()
lc += Loss_cot.item()
ld += Loss_diff.item()
# using tensorboard to monitor loss and acc
writer.add_scalars(
'data/loss', {
'loss_sup': Loss_sup.item(),
'loss_cot': Loss_cot.item(),
'loss_diff': Loss_diff.item()
}, (epoch) * (step) + i)
writer.add_scalars(
'data/training_accuracy', {
'net1 acc': 100. * (train_correct_S1) / (total_S1),
'net2 acc': 100. * (train_correct_S2) / (total_S2)
}, (epoch) * (step) + i)
if (i + 1) % 50 == 0:
# print statistics
tqdm.write(
'net1 training acc: %.3f%% | net2 training acc: %.3f%% | total loss: %.3f | loss_sup: %.3f | loss_cot: %.3f | loss_diff: %.3f '
% (100. * (train_correct_S1 + train_correct_U1) /
(total_S1 + total_U1), 100. *
(train_correct_S2 + train_correct_U2) /
(total_S2 + total_U2), running_loss / (i + 1), ls /
(i + 1), lc / (i + 1), ld / (i + 1)))
#attack
"""
from advertorch.attacks import LinfPGDAttack
adversary = LinfPGDAttack(model, loss_fn=nn.CrossEntropyLoss(reduction="sum"), eps=0.15,
nb_iter=40, eps_iter=0.01, rand_init=True, clip_min=0.0, clip_max=1.0,
targeted=False)
"""
from advertorch.attacks import GradientSignAttack
import time
t = time.time()
adversary = GradientSignAttack(
model,
loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=0.02,
clip_min=0.0,
clip_max=1.0,
targeted=False)
adv_untarget = adversary.perturb(cln_data, true_label)
target = torch.ones_like(true_label) * 3
adversary.targeted = True
adv_targeted = adversary.perturb(cln_data, target)
pred_cln = get_cls(model(cln_data))
pred_untarget = get_cls(model(adv_untarget))
pred_target = get_cls(model(adv_targeted))
import matplotlib.pyplot as plt
plt.figure(figsize=(10, 8))
def run(cfg: DictConfig) -> DictConfig:
# keep the file directory as the current working directory
os.chdir(hydra.utils.get_original_cwd())
reset_seed(cfg.train_param.seed)
# -------- Get the pre-processor --------
train_transform, val_transform = load_preprocesser(cfg.dataset.dataset, "dct")
# -------- Get the dataset --------
manager, train_loader, val_loader = load_dataset(
cfg.dataset.dataset,
"dct",
dataset_root=cfg.path.dataset_root,
supervised_ratio=cfg.train_param.supervised_ratio,
batch_size=cfg.train_param.batch_size,
train_folds=cfg.train_param.train_folds,
val_folds=cfg.train_param.val_folds,
train_transform=train_transform,
val_transform=val_transform,
num_workers=0,
pin_memory=True,
verbose=1
)
# The input shape needed to generate the model
input_shape = train_loader._iterables[0].dataset[0][0].shape
# -------- Prepare the models ---------
torch.cuda.empty_cache()
model_func = load_model(cfg.dataset.dataset, cfg.model.model)
commun_args = dict(
manager=manager,
num_classes=cfg.dataset.num_classes,
input_shape=list(input_shape),
)
m1, m2 = model_func(**commun_args), model_func(**commun_args)
m1, m2 = m1.cuda(), m2.cuda()
summary(m1, input_shape)
# --------- Tensorboard and Checkpoint title ---------
# Prepare sufix
sufix_title = ''
sufix_title += f'_{cfg.train_param.learning_rate}-lr'
sufix_title += f'_{cfg.train_param.supervised_ratio}-sr'
sufix_title += f'_{cfg.train_param.nb_epoch}-e'
sufix_title += f'_{cfg.train_param.batch_size}-bs'
sufix_title += f'_{cfg.train_param.seed}-seed'
# deep-co-training parameters
sufix_title = f'_{cfg.dct.epsilon}-eps'
sufix_title += f'_{cfg.dct.warmup_length}-wl'
sufix_title += f'_{cfg.dct.lambda_cot_max}-lcm'
sufix_title += f'_{cfg.dct.lambda_diff_max}-ldm'
# -------- Tensorboard logging --------
tensorboard_title = f'{get_datetime()}_{cfg.model.model}_{sufix_title}'
log_dir = f'{cfg.path.tensorboard_path}/{tensorboard_title}'
print('Tensorboard log at: ', log_dir)
tensorboard = mSummaryWriter(log_dir=log_dir, comment=cfg.model.model)
# --------- Optimizer, callbacks, loss and checkpoint ---------
optimizer = load_optimizer(cfg.dataset.dataset, "dct",
model1=m1, model2=m2, learning_rate=cfg.train_param.learning_rate)
callbacks = load_callbacks(cfg.dataset.dataset, "dct",
optimizer=optimizer, nb_epoch=cfg.train_param.nb_epoch)
# Losses
# All losses are define in DCT/loss.py
# warmup
lambda_cot = Warmup(cfg.dct.lambda_cot_max, cfg.dct.warmup_length, sigmoid_rampup)
lambda_diff = Warmup(cfg.dct.lambda_diff_max, cfg.dct.warmup_length, sigmoid_rampup)
callbacks += [lambda_cot, lambda_diff]
# Checkpoint
checkpoint_title = f'{cfg.model.model}_{sufix_title}'
checkpoint_path = f'{cfg.path.checkpoint_path}/{checkpoint_title}'
checkpoint = CheckPoint([m1, m2], optimizer, mode="max", name=checkpoint_path)
# --------- Adversarial generation ---------
adv_generator_1 = GradientSignAttack(
m1, loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=cfg.dct.epsilon, clip_min=-np.inf, clip_max=np.inf, targeted=False
)
adv_generator_2 = GradientSignAttack(
m2, loss_fn=nn.CrossEntropyLoss(reduction="sum"),
eps=cfg.dct.epsilon, clip_min=-np.inf, clip_max=np.inf, targeted=False
)
# --------- Metric and print formater ---------
metrics_fn = dict(
ratio_s=[Ratio(), Ratio()],
ratio_u=[Ratio(), Ratio()],
acc_s=[CategoricalAccuracy(), CategoricalAccuracy()],
acc_u=[CategoricalAccuracy(), CategoricalAccuracy()],
f1_s=[FScore(), FScore()],
f1_u=[FScore(), FScore()],
avg_total=ContinueAverage(),
avg_sup=ContinueAverage(),
avg_cot=ContinueAverage(),
avg_diff=ContinueAverage(),
)
maximum_tracker = track_maximum()
def reset_metrics():
for item in metrics_fn.values():
if isinstance(item, list):
for f in item:
f.reset()
else:
item.reset()
header, train_formater, val_formater = get_train_format('dct')
# --------- Training and Validation function ---------
def train(epoch):
start_time = time.time()
print("")
reset_metrics()
m1.train()
m2.train()
for batch, (S1, S2, U) in enumerate(train_loader):
x_s1, y_s1 = S1
x_s2, y_s2 = S2
x_u, y_u = U
x_s1, x_s2, x_u = x_s1.cuda().float(), x_s2.cuda().float(), x_u.cuda().float()
y_s1, y_s2, y_u = y_s1.cuda().long(), y_s2.cuda().long(), y_u.cuda().long()
logits_s1 = m1(x_s1)
logits_s2 = m2(x_s2)
logits_u1 = m1(x_u)
logits_u2 = m2(x_u)
# pseudo labels of U
pred_u1 = torch.argmax(logits_u1, 1)
pred_u2 = torch.argmax(logits_u2, 1)
# ======== Generate adversarial examples ========
# fix batchnorm ----
m1.eval()
m2.eval()
# generate adversarial examples ----
adv_data_s1 = adv_generator_1.perturb(x_s1, y_s1)
adv_data_u1 = adv_generator_1.perturb(x_u, pred_u1)
adv_data_s2 = adv_generator_2.perturb(x_s2, y_s2)
adv_data_u2 = adv_generator_2.perturb(x_u, pred_u2)
m1.train()
m2.train()
# predict adversarial examples ----
adv_logits_s1 = m1(adv_data_s2)
adv_logits_s2 = m2(adv_data_s1)
adv_logits_u1 = m1(adv_data_u2)
adv_logits_u2 = m2(adv_data_u1)
# ======== calculate the differents loss ========
# zero the parameter gradients ----
optimizer.zero_grad()
m1.zero_grad()
m2.zero_grad()
# losses ----
l_sup = loss_sup(logits_s1, logits_s2, y_s1, y_s2)
l_cot = loss_cot(logits_u1, logits_u2)
l_diff = loss_diff(
logits_s1, logits_s2, adv_logits_s1, adv_logits_s2,
logits_u1, logits_u2, adv_logits_u1, adv_logits_u2
)
total_loss = l_sup + lambda_cot() * l_cot + lambda_diff() * l_diff
total_loss.backward()
optimizer.step()
# ======== Calc the metrics ========
with torch.set_grad_enabled(False):
# accuracies ----
pred_s1 = torch.argmax(logits_s1, dim=1)
pred_s2 = torch.argmax(logits_s2, dim=1)
acc_s1 = metrics_fn["acc_s"][0](pred_s1, y_s1)
acc_s2 = metrics_fn["acc_s"][1](pred_s2, y_s2)
acc_u1 = metrics_fn["acc_u"][0](pred_u1, y_u)
acc_u2 = metrics_fn["acc_u"][1](pred_u2, y_u)
# ratios ----
adv_pred_s1 = torch.argmax(adv_logits_s1, 1)
adv_pred_s2 = torch.argmax(adv_logits_s2, 1)
adv_pred_u1 = torch.argmax(adv_logits_u1, 1)
adv_pred_u2 = torch.argmax(adv_logits_u2, 1)
ratio_s1 = metrics_fn["ratio_s"][0](adv_pred_s1, y_s1)
ratio_s2 = metrics_fn["ratio_s"][0](adv_pred_s2, y_s2)
ratio_u1 = metrics_fn["ratio_s"][0](adv_pred_u1, y_u)
ratio_u2 = metrics_fn["ratio_s"][0](adv_pred_u2, y_u)
# ========
avg_total = metrics_fn["avg_total"](total_loss.item())
avg_sup = metrics_fn["avg_sup"](l_sup.item())
avg_diff = metrics_fn["avg_diff"](l_diff.item())
avg_cot = metrics_fn["avg_cot"](l_cot.item())
# logs
print(train_formater.format(
"Training: ",
epoch + 1,
int(100 * (batch + 1) / len(train_loader)),
"", avg_sup.mean(size=100), avg_cot.mean(size=100), avg_diff.mean(size=100), avg_total.mean(size=100),
"", acc_s1.mean(size=100), acc_u1.mean(size=100),
time.time() - start_time
), end="\r")
# using tensorboard to monitor loss and acc\n",
tensorboard.add_scalar('train/total_loss', avg_total.mean(size=100), epoch)
tensorboard.add_scalar('train/Lsup', avg_sup.mean(size=100), epoch)
tensorboard.add_scalar('train/Lcot', avg_cot.mean(size=100), epoch)
tensorboard.add_scalar('train/Ldiff', avg_diff.mean(size=100), epoch)
tensorboard.add_scalar("train/acc_1", acc_s1.mean(size=100), epoch)
tensorboard.add_scalar("train/acc_2", acc_s2.mean(size=100), epoch)
tensorboard.add_scalar("detail_acc/acc_s1", acc_s1.mean(size=100), epoch)
tensorboard.add_scalar("detail_acc/acc_s2", acc_s2.mean(size=100), epoch)
tensorboard.add_scalar("detail_acc/acc_u1", acc_u1.mean(size=100), epoch)
tensorboard.add_scalar("detail_acc/acc_u2", acc_u2.mean(size=100), epoch)
tensorboard.add_scalar("detail_ratio/ratio_s1", ratio_s1.mean(size=100), epoch)
tensorboard.add_scalar("detail_ratio/ratio_s2", ratio_s2.mean(size=100), epoch)
tensorboard.add_scalar("detail_ratio/ratio_u1", ratio_u1.mean(size=100), epoch)
tensorboard.add_scalar("detail_ratio/ratio_u2", ratio_u2.mean(size=100), epoch)
# Return the total loss to check for NaN
return total_loss.item()
def val(epoch, msg=""):
start_time = time.time()
print("")
reset_metrics()
m1.eval()
m2.eval()
with torch.set_grad_enabled(False):
for batch, (X, y) in enumerate(val_loader):
x = X.cuda().float()
y = y.cuda().long()
with autocast():
logits_1 = m1(x)
logits_2 = m2(x)
# losses ----
l_sup = loss_sup(logits_1, logits_2, y, y)
# ======== Calc the metrics ========
# accuracies ----
pred_1 = torch.argmax(logits_1, dim=1)
pred_2 = torch.argmax(logits_2, dim=1)
acc_1 = metrics_fn["acc_s"][0](pred_1, y)
acc_2 = metrics_fn["acc_s"][1](pred_2, y)
avg_sup = metrics_fn["avg_sup"](l_sup.item())
# logs
print(val_formater.format(
"Validation: ",
epoch + 1,
int(100 * (batch + 1) / len(train_loader)),
"", avg_sup.mean(size=100), 0.0, 0.0, avg_sup.mean(size=100),
"", acc_1.mean(size=100), 0.0,
time.time() - start_time
), end="\r")
tensorboard.add_scalar("val/acc_1", acc_1.mean(size=100), epoch)
tensorboard.add_scalar("val/acc_2", acc_2.mean(size=100), epoch)
tensorboard.add_scalar(
"max/acc_1", maximum_tracker("acc_1", acc_1.mean(size=100)), epoch)
tensorboard.add_scalar(
"max/acc_2", maximum_tracker("acc_2", acc_2.mean(size=100)), epoch)
tensorboard.add_scalar(
"detail_hyperparameters/lambda_cot", lambda_cot(), epoch)
tensorboard.add_scalar(
"detail_hyperparameters/lambda_diff", lambda_diff(), epoch)
tensorboard.add_scalar(
"detail_hyperparameters/learning_rate", get_lr(optimizer), epoch)
# Apply callbacks
for c in callbacks:
c.step()
# call checkpoint
checkpoint.step(acc_1.mean(size=100))
# -------- Training loop --------
print(header)
if cfg.train_param.resume:
checkpoint.load_last()
start_epoch = checkpoint.epoch_counter
end_epoch = cfg.train_param.nb_epoch
for e in range(start_epoch, end_epoch):
train(e)
val(e)
tensorboard.flush()
# -------- Save the hyper parameter and the metrics
hparams = {
'dataset': cfg.dataset.dataset,
'model': cfg.model.model,
'supervised_ratio': cfg.train_param.supervised_ratio,
'batch_size': cfg.train_param.batch_size,
'nb_epoch': cfg.train_param.nb_epoch,
'learning_rate': cfg.train_param.learning_rate,
'seed': cfg.train_param.seed,
'train_folds': cfg.train_param.train_folds,
'val_folds': cfg.train_param.val_folds,
'epsilon': cfg.dct.epsilon,
'warmup_length': cfg.dct.warmup_length,
'lambda_cot_max': cfg.dct.lambda_cot_max,
'lambda_diff_max': cfg.dct.lambda_diff_max,
}
# convert all value to str
hparams = dict(zip(hparams.keys(), map(str, hparams.values())))
final_metrics = {
"max_acc_1": maximum_tracker.max["acc_1"],
"max_acc_2": maximum_tracker.max["acc_2"],
}
tensorboard.add_hparams(hparams, final_metrics)
tensorboard.flush()
tensorboard.close()
