import datetime
import logging
import math
import time
import torch
import os
from os import path as osp
from pyiqa.data import build_dataloader, build_dataset
from pyiqa.data.data_sampler import EnlargedSampler
from pyiqa.data.prefetch_dataloader import CPUPrefetcher, CUDAPrefetcher
from pyiqa.models import build_model
from pyiqa.utils import (
AvgTimer,
MessageLogger,
check_resume,
get_env_info,
get_root_logger,
get_time_str,
init_tb_logger,
init_wandb_logger,
make_exp_dirs,
mkdir_and_rename,
scandir,
)
from pyiqa.utils.options import copy_opt_file, dict2str, parse_options
from pyiqa.utils.dist_util import master_only
@master_only
[docs]
def init_tb_loggers(opt):
# initialize wandb logger before tensorboard logger to allow proper sync
if (
(opt['logger'].get('wandb') is not None)
and (opt['logger']['wandb'].get('project') is not None)
and ('debug' not in opt['name'])
):
assert opt['logger'].get('use_tb_logger') is True, (
'should turn on tensorboard when using wandb'
)
init_wandb_logger(opt)
tb_logger = None
if opt['logger'].get('use_tb_logger') and 'debug' not in opt['name']:
tb_logger = init_tb_logger(
log_dir=osp.join(opt['root_path'], 'tb_logger', opt['name'])
)
return tb_logger
[docs]
def create_train_val_dataloader(opt, logger):
# create train and val dataloaders
train_loader, val_loaders = None, []
for phase, dataset_opt in opt['datasets'].items():
if phase == 'train':
dataset_enlarge_ratio = dataset_opt.get('dataset_enlarge_ratio', 1)
train_set = build_dataset(dataset_opt)
train_sampler = EnlargedSampler(
train_set,
opt['world_size'],
opt['rank'],
dataset_enlarge_ratio,
dataset_opt.get('use_shuffle', True),
)
train_loader = build_dataloader(
train_set,
dataset_opt,
num_gpu=opt['num_gpu'],
dist=opt['dist'],
sampler=train_sampler,
seed=opt['manual_seed'],
)
num_iter_per_epoch = math.ceil(
len(train_set)
* dataset_enlarge_ratio
/ (dataset_opt['batch_size_per_gpu'] * opt['world_size'])
)
total_epochs = opt['train'].get('total_epoch', None)
if total_epochs is not None:
total_epochs = int(total_epochs)
total_iters = total_epochs * (num_iter_per_epoch)
opt['train']['total_iter'] = total_iters
else:
total_iters = int(opt['train']['total_iter'])
total_epochs = math.ceil(total_iters / (num_iter_per_epoch))
logger.info(
'Training statistics:'
f'\n\tNumber of train images: {len(train_set)}'
f'\n\tDataset enlarge ratio: {dataset_enlarge_ratio}'
f'\n\tBatch size per gpu: {dataset_opt["batch_size_per_gpu"]}'
f'\n\tWorld size (gpu number): {opt["world_size"]}'
f'\n\tRequire iter number per epoch: {num_iter_per_epoch}'
f'\n\tTotal epochs: {total_epochs}; iters: {total_iters}.'
)
elif phase.split('_')[0] == 'val':
val_set = build_dataset(dataset_opt)
val_loader = build_dataloader(
val_set,
dataset_opt,
num_gpu=opt['num_gpu'],
dist=opt['dist'],
sampler=None,
seed=opt['manual_seed'],
)
logger.info(
f'Number of val images/folders in {dataset_opt["name"]}: {len(val_set)}'
)
val_loaders.append(val_loader)
else:
raise ValueError(f'Dataset phase {phase} is not recognized.')
return train_loader, train_sampler, val_loaders, total_epochs, total_iters
[docs]
def load_resume_state(opt):
resume_state_path = None
if opt['auto_resume']:
state_path = osp.join('experiments', opt['name'], 'training_states')
if osp.isdir(state_path):
states = list(
scandir(state_path, suffix='state', recursive=False, full_path=False)
)
if len(states) != 0:
states = [float(v.split('.state')[0]) for v in states]
resume_state_path = osp.join(state_path, f'{max(states):.0f}.state')
opt['path']['resume_state'] = resume_state_path
else:
if opt['path'].get('resume_state'):
resume_state_path = opt['path']['resume_state']
if resume_state_path is None:
resume_state = None
else:
device_id = torch.cuda.current_device()
resume_state = torch.load(
resume_state_path, map_location=lambda storage, loc: storage.cuda(device_id)
)
check_resume(opt, resume_state['iter'])
return resume_state
[docs]
def train_pipeline(root_path, opt=None, args=None):
# parse options, set distributed setting, set random seed
if opt is None and args is None:
opt, args = parse_options(root_path, is_train=True)
opt['root_path'] = root_path
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
# load resume states if necessary
resume_state = load_resume_state(opt)
# mkdir for experiments and logger
if resume_state is None:
make_exp_dirs(opt)
if (
opt['logger'].get('use_tb_logger')
and 'debug' not in opt['name']
and opt['rank'] == 0
):
os.makedirs(osp.join(opt['root_path'], 'tb_logger_archived'), exist_ok=True)
mkdir_and_rename(osp.join(opt['root_path'], 'tb_logger', opt['name']))
# copy the yml file to the experiment root
copy_opt_file(args.opt, opt['path']['experiments_root'])
# WARNING: should not use get_root_logger in the above codes, including the called functions
# Otherwise the logger will not be properly initialized
log_file = osp.join(opt['path']['log'], f'train_{opt["name"]}_{get_time_str()}.log')
logger = get_root_logger(
logger_name='pyiqa', log_level=logging.INFO, log_file=log_file
)
logger.info(get_env_info())
logger.info(dict2str(opt))
# initialize wandb and tb loggers
tb_logger = init_tb_loggers(opt)
# create train and validation dataloaders
result = create_train_val_dataloader(opt, logger)
train_loader, train_sampler, val_loaders, total_epochs, total_iters = result
# create model
model = build_model(opt)
if resume_state: # resume training
model.resume_training(resume_state) # handle optimizers and schedulers
logger.info(
f'Resuming training from epoch: {resume_state["epoch"]}, '
f'iter: {resume_state["iter"]}.'
)
start_epoch = resume_state['epoch']
current_iter = resume_state['iter']
else:
start_epoch = 0
current_iter = 0
# create message logger (formatted outputs)
msg_logger = MessageLogger(opt, current_iter, tb_logger)
# dataloader prefetcher
prefetch_mode = opt['datasets']['train'].get('prefetch_mode')
if prefetch_mode is None or prefetch_mode == 'cpu':
prefetcher = CPUPrefetcher(train_loader)
elif prefetch_mode == 'cuda':
prefetcher = CUDAPrefetcher(train_loader, opt)
logger.info(f'Use {prefetch_mode} prefetch dataloader')
if opt['datasets']['train'].get('pin_memory') is not True:
raise ValueError('Please set pin_memory=True for CUDAPrefetcher.')
else:
raise ValueError(
f'Wrong prefetch_mode {prefetch_mode}.'
"Supported ones are: None, 'cuda', 'cpu'."
)
# training
logger.info(f'Start training from epoch: {start_epoch}, iter: {current_iter}')
data_timer, iter_timer = AvgTimer(), AvgTimer()
start_time = time.time()
for epoch in range(start_epoch, total_epochs + 1):
train_sampler.set_epoch(epoch)
prefetcher.reset()
train_data = prefetcher.next()
while train_data is not None:
data_timer.record()
current_iter += 1
if current_iter > total_iters:
break
# update learning rate
# model.update_learning_rate(current_iter, warmup_iter=opt['train'].get('warmup_iter', -1))
# training
model.feed_data(train_data)
model.optimize_parameters(current_iter)
iter_timer.record()
if current_iter == 1:
# reset start time in msg_logger for more accurate eta_time
# not work in resume mode
msg_logger.reset_start_time()
# log
if current_iter % opt['logger']['print_freq'] == 0:
log_vars = {'epoch': epoch, 'iter': current_iter}
log_vars.update({'lrs': model.get_current_learning_rate()})
log_vars.update(
{
'time': iter_timer.get_avg_time(),
'data_time': data_timer.get_avg_time(),
}
)
log_vars.update(model.get_current_log())
msg_logger(log_vars)
# log images
log_img_freq = opt['logger'].get('log_imgs_freq', 1e99)
if current_iter % log_img_freq == 0:
visual_imgs = model.get_current_visuals()
if tb_logger and visual_imgs is not None:
for k, v in visual_imgs.items():
tb_logger.add_images(
f'ckpt_imgs/{k}', v.clamp(0, 1), current_iter
)
# save models and training states
save_ckpt_freq = opt['logger'].get('save_checkpoint_freq', 9e9)
if current_iter % save_ckpt_freq == 0:
logger.info('Saving models and training states.')
model.save(epoch, current_iter)
if current_iter % opt['logger']['save_latest_freq'] == 0:
logger.info('Saving latest models and training states.')
model.save(epoch, -1)
# validation
if opt.get('val') is not None and (
current_iter % opt['val']['val_freq'] == 0
):
logger.info(
f'{len(val_loaders)} validation datasets are used for validation.'
)
for val_loader in val_loaders:
model.validation(
val_loader, current_iter, tb_logger, opt['val']['save_img']
)
data_timer.start()
iter_timer.start()
train_data = prefetcher.next()
if 'debug' in opt['name'] and current_iter >= 8:
break
# end of iter
# use epoch based learning rate scheduler
model.update_learning_rate(
epoch + 2, warmup_iter=opt['train'].get('warmup_iter', -1)
)
if 'debug' in opt['name'] and epoch >= 2:
break
# end of epoch
consumed_time = str(datetime.timedelta(seconds=int(time.time() - start_time)))
logger.info(f'End of training. Time consumed: {consumed_time}')
logger.info('Save the latest model.')
model.save(epoch=-1, current_iter=-1) # -1 stands for the latest
if opt.get('val') is not None:
for val_loader in val_loaders:
model.validation(
val_loader, current_iter, tb_logger, opt['val']['save_img']
)
if tb_logger:
tb_logger.close()
if opt['rank'] == 0:
return model.best_metric_results
if __name__ == '__main__':
[docs]
root_path = osp.abspath(osp.join(__file__, osp.pardir, osp.pardir))
train_pipeline(root_path)