def main_worker(cfg):
# create tensorboard and logs
if cfg.DDP_CONFIG.GPU_WORLD_RANK == 0:
tb_logdir = build_log_dir(cfg)
writer = SummaryWriter(log_dir=tb_logdir)
else:
writer = None
cfg.freeze()
# create model
model = get_model(cfg)
model = deploy_model(model, cfg)
# create dataset and dataloader
val_loader, _ = build_dataloader_val(cfg)
if cfg.CONFIG.MODEL.LOAD:
model, _ = load_model(model, cfg, load_fc=True)
criterion = nn.CrossEntropyLoss().cuda()
# adversarial_classification(model, val_loader, -1, criterion, cfg, writer)
validation_classification(model, val_loader, -1, criterion, cfg, writer)
if writer is not None:
writer.close()
def main_worker(cfg):
# create tensorboard and logs
if cfg.DDP_CONFIG.GPU_WORLD_RANK == 0:
tb_logdir = build_log_dir(cfg)
writer = SummaryWriter(log_dir=tb_logdir)
else:
writer = None
cfg.freeze()
# create model
model = get_model(cfg)
model = deploy_model(model, cfg)
# create dataset and dataloader
train_loader, val_loader, train_sampler, val_sampler, mg_sampler = build_dataloader(
cfg)
optimizer = torch.optim.SGD(model.parameters(),
lr=cfg.CONFIG.TRAIN.LR,
momentum=cfg.CONFIG.TRAIN.MOMENTUM,
weight_decay=cfg.CONFIG.TRAIN.W_DECAY)
if cfg.CONFIG.MODEL.LOAD:
model, _ = load_model(model, optimizer, cfg, load_fc=True)
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=cfg.CONFIG.TRAIN.LR_MILESTONE,
gamma=cfg.CONFIG.TRAIN.STEP)
criterion = nn.CrossEntropyLoss().cuda()
base_iter = 0
for epoch in range(cfg.CONFIG.TRAIN.EPOCH_NUM):
if cfg.DDP_CONFIG.DISTRIBUTED:
train_sampler.set_epoch(epoch)
base_iter = train_classification(base_iter,
model,
train_loader,
epoch,
criterion,
optimizer,
cfg,
writer=writer)
scheduler.step()
if epoch % cfg.CONFIG.VAL.FREQ == 0 or epoch == cfg.CONFIG.TRAIN.EPOCH_NUM - 1:
validation_classification(model, val_loader, epoch, criterion, cfg,
writer)
if epoch % cfg.CONFIG.LOG.SAVE_FREQ == 0:
if cfg.DDP_CONFIG.GPU_WORLD_RANK == 0 or cfg.DDP_CONFIG.DISTRIBUTED == False:
save_model(model, optimizer, epoch, cfg)
if writer is not None:
writer.close()
def main_worker(cfg):
# create tensorboard and logs
if cfg.DDP_CONFIG.GPU_WORLD_RANK == 0:
tb_logdir = build_log_dir(cfg)
writer = SummaryWriter(log_dir=tb_logdir)
else:
writer = None
cfg.freeze()
# create model
model = get_model(cfg)
model = deploy_model(model, cfg)
# create dataset and dataloader
train_loader, val_loader, train_sampler, val_sampler, mg_sampler = build_dataloader(
cfg)
optimizer = torch.optim.SGD(model.parameters(),
lr=cfg.CONFIG.TRAIN.LR,
momentum=cfg.CONFIG.TRAIN.MOMENTUM,
weight_decay=cfg.CONFIG.TRAIN.W_DECAY)
if cfg.CONFIG.MODEL.LOAD:
model, _ = load_model(model, optimizer, cfg, load_fc=True)
if cfg.CONFIG.TRAIN.LR_POLICY == 'Step':
scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer,
milestones=cfg.CONFIG.TRAIN.LR_MILESTONE,
gamma=cfg.CONFIG.TRAIN.STEP)
elif cfg.CONFIG.TRAIN.LR_POLICY == 'Cosine':
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer,
T_max=cfg.CONFIG.TRAIN.EPOCH_NUM - cfg.CONFIG.TRAIN.WARMUP_EPOCHS,
eta_min=0,
last_epoch=cfg.CONFIG.TRAIN.RESUME_EPOCH)
else:
print(
'Learning rate schedule %s is not supported yet. Please use Step or Cosine.'
)
if cfg.CONFIG.TRAIN.USE_WARMUP:
scheduler_warmup = GradualWarmupScheduler(
optimizer,
multiplier=(cfg.CONFIG.TRAIN.WARMUP_END_LR / cfg.CONFIG.TRAIN.LR),
total_epoch=cfg.CONFIG.TRAIN.WARMUP_EPOCHS,
after_scheduler=scheduler)
criterion = nn.CrossEntropyLoss().cuda()
base_iter = 0
for epoch in range(cfg.CONFIG.TRAIN.EPOCH_NUM):
if cfg.DDP_CONFIG.DISTRIBUTED:
train_sampler.set_epoch(epoch)
base_iter = train_classification(base_iter,
model,
train_loader,
epoch,
criterion,
optimizer,
cfg,
writer=writer)
if cfg.CONFIG.TRAIN.USE_WARMUP:
scheduler_warmup.step()
else:
scheduler.step()
if cfg.CONFIG.TRAIN.MULTIGRID.USE_LONG_CYCLE:
if epoch in cfg.CONFIG.TRAIN.MULTIGRID.LONG_CYCLE_EPOCH:
mg_sampler.step_long_cycle()
if epoch % cfg.CONFIG.VAL.FREQ == 0 or epoch == cfg.CONFIG.TRAIN.EPOCH_NUM - 1:
validation_classification(model, val_loader, epoch, criterion, cfg,
writer)
if epoch % cfg.CONFIG.LOG.SAVE_FREQ == 0:
if cfg.DDP_CONFIG.GPU_WORLD_RANK == 0 or cfg.DDP_CONFIG.DISTRIBUTED == False:
save_model(model, optimizer, epoch, cfg)
if writer is not None:
writer.close()