def main():
global global_step
config = load_config()
set_seed(config)
setup_cudnn(config)
epoch_seeds = np.random.randint(np.iinfo(np.int32).max // 2,
size=config.scheduler.epochs)
if config.train.distributed:
dist.init_process_group(backend=config.train.dist.backend,
init_method=config.train.dist.init_method,
rank=config.train.dist.node_rank,
world_size=config.train.dist.world_size)
torch.cuda.set_device(config.train.dist.local_rank)
output_dir = pathlib.Path(config.train.output_dir)
if get_rank() == 0:
if not config.train.resume and output_dir.exists():
raise RuntimeError(
f'Output directory `{output_dir.as_posix()}` already exists')
output_dir.mkdir(exist_ok=True, parents=True)
if not config.train.resume:
save_config(config, output_dir / 'config.yaml')
save_config(get_env_info(config), output_dir / 'env.yaml')
diff = find_config_diff(config)
if diff is not None:
save_config(diff, output_dir / 'config_min.yaml')
logger = create_logger(name=__name__,
distributed_rank=get_rank(),
output_dir=output_dir,
filename='log.txt')
logger.info(config)
logger.info(get_env_info(config))
train_loader, val_loader = create_dataloader(config, is_train=True)
model = create_model(config)
macs, n_params = count_op(config, model)
logger.info(f'MACs : {macs}')
logger.info(f'#params: {n_params}')
optimizer = create_optimizer(config, model)
model, optimizer = apex.amp.initialize(model,
optimizer,
opt_level=config.train.precision)
model = apply_data_parallel_wrapper(config, model)
scheduler = create_scheduler(config,
optimizer,
steps_per_epoch=len(train_loader))
checkpointer = Checkpointer(model,
optimizer=optimizer,
scheduler=scheduler,
save_dir=output_dir,
save_to_disk=get_rank() == 0)
start_epoch = config.train.start_epoch
scheduler.last_epoch = start_epoch
if config.train.resume:
checkpoint_config = checkpointer.resume_or_load('', resume=True)
global_step = checkpoint_config['global_step']
start_epoch = checkpoint_config['epoch']
config.defrost()
config.merge_from_other_cfg(ConfigNode(checkpoint_config['config']))
config.freeze()
elif config.train.checkpoint != '':
checkpoint = torch.load(config.train.checkpoint, map_location='cpu')
if isinstance(model,
(nn.DataParallel, nn.parallel.DistributedDataParallel)):
model.module.load_state_dict(checkpoint['model'])
else:
model.load_state_dict(checkpoint['model'])
if get_rank() == 0 and config.train.use_tensorboard:
tensorboard_writer = create_tensorboard_writer(
config, output_dir, purge_step=config.train.start_epoch + 1)
tensorboard_writer2 = create_tensorboard_writer(
config, output_dir / 'running', purge_step=global_step + 1)
else:
tensorboard_writer = DummyWriter()
tensorboard_writer2 = DummyWriter()
train_loss, val_loss = create_loss(config)
if (config.train.val_period > 0 and start_epoch == 0
and config.train.val_first):
validate(0, config, model, val_loss, val_loader, logger,
tensorboard_writer)
for epoch, seed in enumerate(epoch_seeds[start_epoch:], start_epoch):
epoch += 1
np.random.seed(seed)
train(epoch, config, model, optimizer, scheduler, train_loss,
train_loader, logger, tensorboard_writer, tensorboard_writer2)
if config.train.val_period > 0 and (epoch %
config.train.val_period == 0):
validate(epoch, config, model, val_loss, val_loader, logger,
tensorboard_writer)
tensorboard_writer.flush()
tensorboard_writer2.flush()
if (epoch % config.train.checkpoint_period == 0) or (
epoch == config.scheduler.epochs):
checkpoint_config = {
'epoch': epoch,
'global_step': global_step,
'config': config.as_dict(),
}
checkpointer.save(f'checkpoint_{epoch:05d}', **checkpoint_config)
tensorboard_writer.close()
tensorboard_writer2.close()
def main():
global global_step
config = load_config()
set_seed(config)
setup_cudnn(config)
# np.iinfo(np_type).max: machine limit (upper bound) of the this type
# every epoch will have a specific epoch seed
epoch_seeds = np.random.randint(np.iinfo(np.int32).max // 2,
size=config.scheduler.epochs)
if config.train.distributed:
dist.init_process_group(backend=config.train.dist.backend,
init_method=config.train.dist.init_method,
rank=config.train.dist.node_rank,
world_size=config.train.dist.world_size)
torch.cuda.set_device(config.train.dist.local_rank)
output_dir = pathlib.Path(config.train.output_dir)
if get_rank() == 0:
if not config.train.resume and output_dir.exists():
raise RuntimeError(
f'Output directory `{output_dir.as_posix()}` already exists')
output_dir.mkdir(exist_ok=True, parents=True)
if not config.train.resume:
# if we need to resume training, current config, environment info and the difference between
# the current and default config will be saved.
save_config(config, output_dir / 'config.yaml')
save_config(get_env_info(config), output_dir / 'env.yaml')
diff = find_config_diff(config)
if diff is not None:
save_config(diff, output_dir / 'config_min.yaml')
logger = create_logger(name=__name__,
distributed_rank=get_rank(),
output_dir=output_dir,
filename='log.txt')
logger.info(config)
logger.info(get_env_info(config))
train_loader, val_loader = create_dataloader(config, is_train=True)
model = create_model(config)
# Multiply-and-ACcumulate(MAC): ops
macs, n_params = count_op(config, model)
logger.info(f'MACs : {macs}')
logger.info(f'#params: {n_params}')
# creating optimizer: SGD with nesterov momentum, adam, amsgrad, adabound, adaboundw or lars.
optimizer = create_optimizer(config, model)
# some AMP(Automatic mixed precision) settings
if config.device != 'cpu':
model, optimizer = apex.amp.initialize(
model, optimizer, opt_level=config.train.precision)
# create data parallel model or distributed data
model = apply_data_parallel_wrapper(config, model)
# set up scheduler and warm up scheduler
# steps per epoch: how many batches in an epoch
scheduler = create_scheduler(config,
optimizer,
steps_per_epoch=len(train_loader))
# create checkponit, do ot use torch's default checkpoint saver because it can't save scheduler
checkpointer = Checkpointer(model,
optimizer=optimizer,
scheduler=scheduler,
save_dir=output_dir,
save_to_disk=get_rank() == 0)
start_epoch = config.train.start_epoch
# last_epoch is used to resume training, here normally we should start from config.train.start_epoch
scheduler.last_epoch = start_epoch
# The resume training supports multiple modes:
# 1. resume = True, loading model from the last training checkpoint and following the global step and config
# 2. resume = False, training checkpoint is specified, load checkpoint to cpu
if config.train.resume:
checkpoint_config = checkpointer.resume_or_load('', resume=True)
global_step = checkpoint_config['global_step']
start_epoch = checkpoint_config['epoch']
config.defrost()
config.merge_from_other_cfg(ConfigNode(checkpoint_config['config']))
config.freeze()
elif config.train.checkpoint != '':
checkpoint = torch.load(config.train.checkpoint, map_location='cpu')
if isinstance(model,
(nn.DataParallel, nn.parallel.DistributedDataParallel)):
model.module.load_state_dict(checkpoint['model'])
else:
model.load_state_dict(checkpoint['model'])
# Two TensorBoard writer:
# First writer for this run of training(maybe it's resuming training)
# Second writer follows the global steps and records the global run.
if get_rank() == 0 and config.train.use_tensorboard:
tensorboard_writer = create_tensorboard_writer(
config, output_dir, purge_step=config.train.start_epoch + 1)
tensorboard_writer2 = create_tensorboard_writer(
config, output_dir / 'running', purge_step=global_step + 1)
else:
tensorboard_writer = DummyWriter()
tensorboard_writer2 = DummyWriter()
train_loss, val_loss = create_loss(config)
if (config.train.val_period > 0 and start_epoch == 0
and config.train.val_first):
# validate the model from epoch 0
validate(0, config, model, val_loss, val_loader, logger,
tensorboard_writer)
for epoch, seed in enumerate(epoch_seeds[start_epoch:], start_epoch):
epoch += 1
np.random.seed(seed)
train(epoch, config, model, optimizer, scheduler, train_loss,
train_loader, logger, tensorboard_writer, tensorboard_writer2)
if config.train.val_period > 0 and (epoch % config.train.val_period
== 0):
validate(epoch, config, model, val_loss, val_loader, logger,
tensorboard_writer)
tensorboard_writer.flush()
tensorboard_writer2.flush()
if (epoch % config.train.checkpoint_period
== 0) or (epoch == config.scheduler.epochs):
checkpoint_config = {
'epoch': epoch,
'global_step': global_step,
'config': config.as_dict(),
}
checkpointer.save(f'checkpoint_{epoch:05d}', **checkpoint_config)
tensorboard_writer.close()
tensorboard_writer2.close()