def train_example(loggers, loaders, model, optimizer, scheduler):
start_epoch = 0
if cfg.train.auto_resume:
start_epoch = load_ckpt(model, optimizer, scheduler,
cfg.train.epoch_resume)
if start_epoch == cfg.optim.max_epoch:
logging.info('Checkpoint found, Task already done')
else:
logging.info('Start from epoch %s', start_epoch)
num_splits = len(loggers)
for cur_epoch in range(start_epoch, cfg.optim.max_epoch):
train_epoch(loggers[0], loaders[0], model, optimizer, scheduler)
loggers[0].write_epoch(cur_epoch)
if is_eval_epoch(cur_epoch):
for i in range(1, num_splits):
eval_epoch(loggers[i], loaders[i], model)
loggers[i].write_epoch(cur_epoch)
if is_ckpt_epoch(cur_epoch):
save_ckpt(model, optimizer, scheduler, cur_epoch)
for logger in loggers:
logger.close()
if cfg.train.ckpt_clean:
clean_ckpt()
logging.info('Task done, results saved in %s', cfg.run_dir)
def train(loggers, loaders, model, optimizer, scheduler):
start_epoch = 0
if cfg.train.auto_resume:
start_epoch = load_ckpt(model, optimizer, scheduler)
if start_epoch == cfg.optim.max_epoch:
logging.info('Checkpoint found, Task already done')
else:
logging.info('Start from epoch {}'.format(start_epoch))
num_splits = len(loggers)
split_names = ['val', 'test']
for cur_epoch in range(start_epoch, cfg.optim.max_epoch):
train_epoch(loggers[0], loaders[0], model, optimizer, scheduler)
loggers[0].write_epoch(cur_epoch)
if is_eval_epoch(cur_epoch):
for i in range(1, num_splits):
eval_epoch(loggers[i],
loaders[i],
model,
split=split_names[i - 1])
loggers[i].write_epoch(cur_epoch)
if is_ckpt_epoch(cur_epoch):
save_ckpt(model, optimizer, scheduler, cur_epoch)
for logger in loggers:
logger.close()
if cfg.train.ckpt_clean:
clean_ckpt()
logging.info('Task done, results saved in {}'.format(cfg.run_dir))
def train(loggers, loaders, model, optimizer, scheduler):
"""
The core training pipeline
Args:
loggers: List of loggers
loaders: List of loaders
model: GNN model
optimizer: PyTorch optimizer
scheduler: PyTorch learning rate scheduler
"""
start_epoch = 0
if cfg.train.auto_resume:
start_epoch = load_ckpt(model, optimizer, scheduler,
cfg.train.epoch_resume)
if start_epoch == cfg.optim.max_epoch:
logging.info('Checkpoint found, Task already done')
else:
logging.info('Start from epoch {}'.format(start_epoch))
num_splits = len(loggers)
split_names = ['val', 'test']
for cur_epoch in range(start_epoch, cfg.optim.max_epoch):
train_epoch(loggers[0], loaders[0], model, optimizer, scheduler)
if is_train_eval_epoch(cur_epoch):
loggers[0].write_epoch(cur_epoch)
if is_eval_epoch(cur_epoch):
for i in range(1, num_splits):
eval_epoch(loggers[i],
loaders[i],
model,
split=split_names[i - 1])
loggers[i].write_epoch(cur_epoch)
if is_ckpt_epoch(cur_epoch) and cfg.train.enable_ckpt:
save_ckpt(model, optimizer, scheduler, cur_epoch)
for logger in loggers:
logger.close()
if cfg.train.ckpt_clean:
clean_ckpt()
logging.info('Task done, results saved in {}'.format(cfg.run_dir))