args.resume = os.path.join(config.model_dir, 'last.checkpoint')
if args.resume is not None and os.path.isfile(args.resume):
now_epoch = load_checkpoint(args.resume, net, optimizer, lr_scheduler)
now_train_time = 0
while True:
if now_epoch > config.num_epochs:
break
now_epoch = now_epoch + 1
descrip_str = 'Training epoch:{}/{} -- lr:{}'.format(
now_epoch, config.num_epochs,
lr_scheduler.get_lr()[0])
s_time = time.time()
acc, yofoacc = train_one_epoch(net, ds_train, optimizer, criterion,
LayerOneTrainer, config.K, DEVICE,
descrip_str)
now_train_time = now_train_time + time.time() - s_time
tb_train_dic = {'Acc': acc, 'YofoAcc': yofoacc}
print(tb_train_dic)
writer.add_scalars('Train', tb_train_dic, now_epoch)
if config.val_interval > 0 and now_epoch % config.val_interval == 0:
acc, advacc = eval_one_epoch(net, ds_val, DEVICE, EvalAttack)
tb_val_dic = {'Acc': acc, 'AdvAcc': advacc}
writer.add_scalars('Val', tb_val_dic, now_epoch)
tb_val_dic['time'] = now_train_time
log_str = json.dumps(tb_val_dic)
with open('time.log', 'a') as f:
f.write(log_str + '\n')
lr_scheduler.step()
ds_train = create_train_dataset(args.batch_size)
ds_val = create_test_dataset(args.batch_size)
TrainAttack = config.create_attack_method(DEVICE)
EvalAttack = config.create_evaluation_attack_method(DEVICE)
now_epoch = 0
if args.auto_continue:
args.resume = os.path.join(config.model_dir, 'last.checkpoint')
if args.resume is not None and os.path.isfile(args.resume):
now_epoch = load_checkpoint(args.resume, net, optimizer,lr_scheduler)
while True:
if now_epoch > config.num_epochs:
break
now_epoch = now_epoch + 1
descrip_str = 'Training epoch:{}/{} -- lr:{}'.format(now_epoch, config.num_epochs,
lr_scheduler.get_lr()[0])
train_one_epoch(net, ds_train, optimizer, criterion, DEVICE,
descrip_str, TrainAttack, config.alpha)
if config.val_interval > 0 and now_epoch % config.val_interval == 0:
eval_one_epoch(net, ds_val, DEVICE, EvalAttack)
lr_scheduler.step()
save_checkpoint(now_epoch, net, optimizer, lr_scheduler,
file_name = os.path.join(config.model_dir, 'epoch-{}.checkpoint'.format(now_epoch)))