lr = optimizer.param_groups[0]['lr']
logger.info(
f'Train: [{epoch}/{args.epochs}][{idx}/{train_len}] '
f'Time {batch_time.val:.3f} ({batch_time.avg:.3f}) '
f'lr {lr:.3f} '
f'loss {loss_meter.val:.3f} ({loss_meter.avg:.3f})')
# tensorboard logger
if summary_writer is not None:
step = (epoch - 1) * len(train_loader) + idx
summary_writer.add_scalar('lr', lr, step)
summary_writer.add_scalar('loss', loss_meter.val, step)
if __name__ == '__main__':
opt = parse_option(stage='pre-train')
if opt.amp_opt_level != "O0":
assert amp is not None, "amp not installed!"
torch.cuda.set_device(opt.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
cudnn.benchmark = True
# setup logger
os.makedirs(opt.output_dir, exist_ok=True)
logger = setup_logger(output=opt.output_dir, distributed_rank=dist.get_rank(), name="contrast")
if dist.get_rank() == 0:
path = os.path.join(opt.output_dir, "config.json")
with open(path, 'w') as f:
json.dump(vars(opt), f, indent=2)
if idx % args.print_freq == 0:
logger.info(
f'Test: [{idx}/{len(val_loader)}]\t'
f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
f'Loss {loss_meter.val:.4f} ({loss_meter.avg:.4f})\t'
f'Acc@1 {acc1_meter.val:.3f} ({acc1_meter.avg:.3f})\t'
f'Acc@5 {acc5_meter.val:.3f} ({acc5_meter.avg:.3f})')
logger.info(
f' * Acc@1 {acc1_meter.avg:.3f} Acc@5 {acc5_meter.avg:.3f}')
return acc1_meter.avg, acc5_meter.avg, loss_meter.avg
if __name__ == '__main__':
opt = parse_option(stage='linear')
if opt.amp_opt_level != "O0":
assert amp is not None, "amp not installed!"
torch.cuda.set_device(opt.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
cudnn.benchmark = True
best_acc1 = 0
os.makedirs(opt.output_dir, exist_ok=True)
logger = setup_logger(output=opt.output_dir,
distributed_rank=dist.get_rank(),
name="contrast")
if dist.get_rank() == 0:
path = os.path.join(opt.output_dir, "config.json")