def get_model(event, context):
"""
Get information about a particular model
Trigger event: GET https://.../models/{model_id}
Inputs:
- path parameter: `model_id` (required)
Returns: An information dictionary about the requested model
"""
model_id = event['pathParameters']['model_id']
return success(model.get_model(model_id))
def generate_adversarial_images(args):
# assertions
assert args.adversary_to_generate is not None, \
"adversary_to_generate can't be None"
assert AdversaryType.has_value(args.adversary_to_generate), \
"\"{}\" adversary_to_generate not defined".format(args.adversary_to_generate)
defense_name = None if not args.defenses else args.defenses[0]
# defense = get_defense(defense_name, args)
data_indices = _get_data_indices(args)
data_type = args.data_type if args.data_type == "train" else "valid"
dataset = load_dataset(args, data_type, None, data_indices=data_indices)
data_loader = get_data_loader(
dataset,
batchsize=args.batchsize,
device=args.device,
shuffle=False)
model, _, _ = get_model(args, load_checkpoint=True, defense_name=defense_name)
adv_params = constants.get_adv_params(args)
print('| adv_params:', adv_params)
status = None
all_inputs = None
all_outputs = None
all_targets = None
bar = progressbar.ProgressBar(len(data_loader))
bar.start()
for batch_num, (imgs, targets) in enumerate(data_loader):
if args.adversary_to_generate == str(AdversaryType.DEEPFOOL):
assert adv_params['learning_rate'] is not None
s, r = adversary.deepfool(
model, imgs, targets, args.data_params['NUM_CLASSES'],
train_mode=(args.data_type == 'train'), max_iter=args.max_adv_iter,
step_size=adv_params['learning_rate'], batch_size=args.batchsize,
labels=dataset.get_classes())
elif args.adversary_to_generate == str(AdversaryType.FGS):
s, r = adversary.fgs(
model, imgs, targets, train_mode=(args.data_type == 'train'),
mode=args.fgs_mode)
elif args.adversary_to_generate == str(AdversaryType.IFGS):
assert adv_params['learning_rate'] is not None
s, r = adversary.ifgs(
model, imgs, targets,
train_mode=(args.data_type == 'train'), max_iter=args.max_adv_iter,
step_size=adv_params['learning_rate'], mode=args.fgs_mode)
elif args.adversary_to_generate == str(AdversaryType.CWL2):
assert args.adv_strength is not None and len(args.adv_strength) == 1
if len(args.crop_frac) == 1:
crop_frac = args.crop_frac[0]
else:
crop_frac = 1.0
s, r = adversary.cw(
model, imgs, targets, args.adv_strength[0], 'l2',
tv_weight=args.tvm_weight,
train_mode=(args.data_type == 'train'), max_iter=args.max_adv_iter,
drop_rate=args.pixel_drop_rate, crop_frac=crop_frac,
kappa=args.margin)
elif args.adversary_to_generate == str(AdversaryType.CWLINF):
assert args.adv_strength is not None and len(args.adv_strength) == 1
s, r = adversary.cw(
model, imgs, targets, args.adv_strength[0], 'linf',
bound=args.adv_bound,
tv_weight=args.tvm_weight,
train_mode=(args.data_type == 'train'), max_iter=args.max_adv_iter,
drop_rate=args.pixel_drop_rate, crop_frac=args.crop_frac,
kappa=args.margin)
if status is None:
status = s.clone()
all_inputs = imgs.clone()
all_outputs = imgs + r
all_targets = targets.clone()
else:
status = torch.cat((status, s), 0)
all_inputs = torch.cat((all_inputs, imgs), 0)
all_outputs = torch.cat((all_outputs, imgs + r), 0)
all_targets = torch.cat((all_targets, targets), 0)
bar.update(batch_num)
print("| computing adversarial stats...")
if args.compute_stats:
rb, ssim, sc = adversary.compute_stats(all_inputs, all_outputs, status)
print('| average robustness = ' + str(rb))
print('| average SSIM = ' + str(ssim))
print('| success rate = ' + str(sc))
# Unnormalize before saving
unnormalize = Unnormalize(args.data_params['MEAN_STD']['MEAN'],
args.data_params['MEAN_STD']['STD'])
all_inputs = unnormalize(all_inputs)
all_outputs = unnormalize(all_outputs)
# save output
output_file = get_adversarial_file_path(
args, args.adversarial_root, defense_name, adv_params,
data_indices['end_idx'], start_idx=data_indices['start_idx'],
with_defense=False)
print("| Saving adversarial data at " + output_file)
if not os.path.isdir(args.adversarial_root):
os.makedirs(args.adversarial_root)
torch.save({'status': status, 'all_inputs': all_inputs,
'all_outputs': all_outputs, 'all_targets': all_targets},
output_file)
def main():
config = get_config()
set_seed(config.train.seed)
message_info(config)
# create log path
val_logFile, writer_logFile, save_path = create_logFile(config)
# dist --init
if config.dist:
dist.init_process_group(backend=config.dist_backend,
init_method=config.dist_init_method)
torch.cuda.set_device(config.dist_local_rank)
# tensorboard
if get_rank() == 0:
writer = get_tensorboard_writer(writer_logFile, purge_step=None)
else:
writer = None
# model
model = get_model(config)
# optimizer
optimizer = get_optimizer(config, model)
if config.apex:
model, optimizer = apex.amp.initialize(model,
optimizer,
opt_level=config.apex_mode)
if config.dist:
if config.dist_sync_bn:
if config.apex:
model = convert_syncbn_model(model)
else:
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
if config.apex:
model = DDP(model, delay_allreduce=True)
else:
model = nn.parallel.DistributedDataParallel(model,
device_ids=[config.dist_local_rank],
output_device=config.dist_local_rank)
# loss
train_loss, val_loss = get_loss(config)
# load_data
data = pd.read_csv(config.train.dataset)
skf = KFold(n_splits=10, shuffle=True, random_state=452)
for fold_idx, (train_idx, val_idx) in enumerate(
skf.split(data['filename'].values, data['filename'].values)):
if fold_idx == config.train.fold:
break
# create dataloader
train_data = data.iloc[train_idx]
val_data = data.iloc[val_idx]
train_loader = create_dataloader(config, train_data, 'train')
val_loader = create_dataloader(config, val_data, 'val')
if get_rank() == 0:
logger.info(f"Splited train set: {train_data.shape}")
logger.info(f"Splited val set: {val_data.shape}")
# scheduler
scheduler = CosineWarmupLr(config, optimizer, len(train_loader))
best_precision, lowest_loss = 0, 100
for epoch in range(config.train.epoches):
# if config.dist:
# train_loader.sampler.set_epoch(epoch)
# train
train(config, epoch, train_loader, model, optimizer, scheduler, train_loss, writer)
# val
if epoch % config.train.val_preiod == 0:
precision, avg_loss = val(config, val_loader, model, val_loss, writer)
if get_rank() == 0:
with open(val_logFile, 'a') as acc_file:
acc_file.write(
f'Fold: {fold_idx:2d}, '
f'Epoch: {epoch:2d}, '
f'Precision: {precision:.8f}, '
f'Loss: {avg_loss:.8f}\n')
is_best = precision > best_precision
is_lowest_loss = avg_loss < lowest_loss
best_precision = max(precision, best_precision)
lowest_loss = min(avg_loss, lowest_loss)
state = {
'epoch': epoch,
'state_dict': model.state_dict(),
'best_precision': best_precision,
'lowest_loss': lowest_loss,
}
save_checkpoint(state, epoch, is_best, is_lowest_loss, save_path)
torch.save(model,
os.path.join(save_path, 'whole_model.pth'))
if get_rank() == 0:
writer.close()
torch.cuda.empty_cache()
def classify_images(args):
# assertions
assert args.ensemble is None or args.ensemble in ENSEMBLE_TYPE, \
"{} not a supported type. Only supported ensembling are {}".format(
args.ensemble, ENSEMBLE_TYPE)
if not args.ensemble:
assert args.ncrops is None or (len(args.ncrops) == 1
and args.ncrops[0] == 1)
if args.defenses is not None:
for d in args.defenses:
assert DefenseType.has_value(d), \
"\"{}\" defense not defined".format(d)
# crops expected for each defense
assert (args.ncrops is None or len(args.ncrops) == len(
args.defenses)), ("Number of crops for each defense is expected")
assert (args.crop_type is None or len(args.crop_type) == len(
args.defenses)), ("crop_type for each defense is expected")
# assert (len(args.crop_frac) == len(args.defenses)), (
# "crop_frac for each defense is expected")
elif args.ncrops is not None:
# no crop ensembling when defense is None
assert len(args.ncrops) == 1
assert args.crop_frac is not None and len(args.crop_frac) == 1, \
"Only one crop_frac is expected as there is no defense"
assert args.crop_type is not None and len(args.crop_type) == 1, \
"Only one crop_type is expected as there is no defense"
if args.defenses is None or len(args.defenses) == 0:
defenses = [None]
else:
defenses = args.defenses
all_defense_probs = None
for idx, defense_name in enumerate(defenses):
# initialize dataset
defense = get_defense(defense_name, args)
# Read preset params for adversary based on args
adv_params = constants.get_adv_params(args, idx)
print("| adv_params: ", adv_params)
# setup crop
ncrops = 1
crop_type = None
crop_frac = 1.0
if args.ncrops:
crop_type = args.crop_type[idx]
crop_frac = args.crop_frac[idx]
if crop_type == 'sliding':
ncrops = 9
else:
ncrops = args.ncrops[idx]
# Init custom crop function
crop = transforms.Crop(crop_type, crop_frac)
# initialize dataset
dataset = load_dataset(args, 'valid', defense, adv_params, crop)
# load model
model, _, _ = get_model(args,
load_checkpoint=True,
defense_name=defense_name)
# get crop probabilities for crops for current defense
probs, targets = _eval_crops(args, dataset, model, defense, crop,
ncrops, crop_type)
if all_defense_probs is None:
all_defense_probs = torch.zeros(len(defenses), len(dataset),
probs.size(2))
# Ensemble crop probabilities
if args.ensemble == 'max':
probs = torch.max(probs, dim=0)[0]
elif args.ensemble == 'avg': # for average ensembling
probs = torch.mean(probs, dim=0)
else: # for no ensembling
assert all_defense_probs.size(0) == 1
probs = probs[0]
all_defense_probs[idx, :, :] = probs
# free memory
dataset = None
model = None
# Ensemble defense probabilities
if args.ensemble == 'max':
all_defense_probs = torch.max(all_defense_probs, dim=0)[0]
elif args.ensemble == 'avg': # for average ensembling
all_defense_probs = torch.mean(all_defense_probs, dim=0)
else: # for no ensembling
assert all_defense_probs.size(0) == 1
all_defense_probs = all_defense_probs[0]
# Calculate top1 and top5 accuracy
prec1, prec5 = accuracy(all_defense_probs, targets, topk=(1, 5))
print('=' * 50)
print('Results for model={}, attack={}, ensemble_type={} '.format(
args.model, args.adversary, args.ensemble))
prec1 = prec1[0]
prec5 = prec5[0]
print('| classification accuracy @1: %2.5f' % (prec1))
print('| classification accuracy @5: %2.5f' % (prec5))
print('| classification error @1: %2.5f' % (100. - prec1))
print('| classification error @5: %2.5f' % (100. - prec5))
print('| done.')
def train_model(args):
# At max 1 defense as no ensembling in training
assert args.defenses is None or len(args.defenses) == 1
defense_name = None if not args.defenses else args.defenses[0]
defense = get_defense(defense_name, args)
# Load model
model, start_epoch, optimizer_ = get_model(args,
load_checkpoint=args.resume,
defense_name=defense_name,
training=True)
# set up optimizer:
optimizer = _get_optimizer(model, args)
# get from checkpoint if available
if start_epoch and optimizer:
args.start_epoch = start_epoch
optimizer.load_state_dict(optimizer_)
# set up criterion:
criterion = nn.CrossEntropyLoss()
if args.device == 'gpu':
# Call .cuda() method on model
criterion = criterion.cuda()
model = model.cuda()
loaders = {}
# set up start-of-epoch hook:
def start_epoch_hook(epoch, model, optimizer):
print('| epoch %d, training:' % epoch)
adjust_learning_rate(args.lr, epoch, optimizer, args.lr_decay,
args.lr_decay_stepsize)
# set up the end-of-epoch hook:
def end_epoch_hook(epoch, model, optimizer, prec1=None, prec5=None):
# print training error:
if prec1 is not None:
print('| training error @1 (epoch %d): %2.5f' %
(epoch, 100. - prec1))
if prec5 is not None:
print('| training error @5 (epoch %d): %2.5f' %
(epoch, 100. - prec5))
# save checkpoint:
print('| epoch %d, testing:' % epoch)
save_checkpoint(
args.models_root, {
'epoch': epoch + 1,
'model_name': args.model,
'model_state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
})
# measure validation error:
prec1, prec5 = test(model, loaders['valid'])
print('| validation error @1 (epoch %d: %2.5f' % (epoch, 100. - prec1))
print('| validation error @5 (epoch %d: %2.5f' % (epoch, 100. - prec5))
def data_loader_hook(epoch):
# Reload data loader for epoch
if args.preprocessed_epoch_data:
print('| epoch %d, Loading data:' % epoch)
for key in {'train', 'valid'}:
# Load validation data only once
if key == 'valid' and 'valid' in loaders:
break
loaders[key] = get_data_loader(
load_dataset(args, key, defense, epoch=epoch),
batchsize=args.batchsize,
device=args.device,
shuffle=True,
)
# if data needs to be loaded only once and is not yet loaded
elif len(loaders) == 0:
print('| epoch %d, Loading data:' % epoch)
for key in {'train', 'valid'}:
loaders[key] = get_data_loader(
load_dataset(args, key, defense),
batchsize=args.batchsize,
device=args.device,
shuffle=True,
)
return loaders['train']
# train the model:
print('| training model...')
train(model,
criterion,
optimizer,
start_epoch_hook=start_epoch_hook,
end_epoch_hook=end_epoch_hook,
data_loader_hook=data_loader_hook,
start_epoch=args.start_epoch,
end_epoch=args.end_epoch,
learning_rate=args.lr)
print('| done.')