import argparse
import logging
import os
import random
import torch
from fastai.callbacks.general_sched import GeneralScheduler, TrainingPhase
from fastai.distributed import *
from fastai.vision import *
from torch.backends import cudnn
from callbacks import DumpPrediction, IterationCallback, TextAccuracy, TopKTextAccuracy
from dataset import ImageDataset, TextDataset
from losses import MultiLosses
from utils import Config, Logger, MyDataParallel, MyConcatDataset
def _set_random_seed(seed):
if seed is not None:
random.seed(seed)
torch.manual_seed(seed)
cudnn.deterministic = True
logging.warning('You have chosen to seed training. '
'This will slow down your training!')
def _get_training_phases(config, n):
lr = np.array(config.optimizer_lr)
periods = config.optimizer_scheduler_periods
sigma = [config.optimizer_scheduler_gamma ** i for i in range(len(periods))]
phases = [TrainingPhase(n * periods[i]).schedule_hp('lr', lr * sigma[i])
for i in range(len(periods))]
return phases
def _get_dataset(ds_type, paths, is_training, config, **kwargs):
kwargs.update({
'img_h': config.dataset_image_height,
'img_w': config.dataset_image_width,
'max_length': config.dataset_max_length,
'case_sensitive': config.dataset_case_sensitive,
'charset_path': config.dataset_charset_path,
'data_aug': config.dataset_data_aug,
'deteriorate_ratio': config.dataset_deteriorate_ratio,
'is_training': is_training,
'multiscales': config.dataset_multiscales,
'one_hot_y': config.dataset_one_hot_y,
})
datasets = [ds_type(p, **kwargs) for p in paths]
if len(datasets) > 1: return MyConcatDataset(datasets)
else: return datasets[0]
def _get_language_databaunch(config):
kwargs = {
'max_length': config.dataset_max_length,
'case_sensitive': config.dataset_case_sensitive,
'charset_path': config.dataset_charset_path,
'smooth_label': config.dataset_smooth_label,
'smooth_factor': config.dataset_smooth_factor,
'one_hot_y': config.dataset_one_hot_y,
'use_sm': config.dataset_use_sm,
}
train_ds = TextDataset(config.dataset_train_roots[0], is_training=True, **kwargs)
valid_ds = TextDataset(config.dataset_test_roots[0], is_training=False, **kwargs)
data = DataBunch.create(
path=train_ds.path,
train_ds=train_ds,
valid_ds=valid_ds,
bs=config.dataset_train_batch_size,
val_bs=config.dataset_test_batch_size,
num_workers=config.dataset_num_workers,
pin_memory=config.dataset_pin_memory)
logging.info(f'{len(data.train_ds)} training items found.')
if not data.empty_val:
logging.info(f'{len(data.valid_ds)} valid items found.')
return data
def _get_databaunch(config):
# An awkward way to reduce loadding data time during test
if config.global_phase == 'test': config.dataset_train_roots = config.dataset_test_roots
train_ds = _get_dataset(ImageDataset, config.dataset_train_roots, True, config)
valid_ds = _get_dataset(ImageDataset, config.dataset_test_roots, False, config)
data = ImageDataBunch.create(
train_ds=train_ds,
valid_ds=valid_ds,
bs=config.dataset_train_batch_size,
val_bs=config.dataset_test_batch_size,
num_workers=config.dataset_num_workers,
pin_memory=config.dataset_pin_memory).normalize(imagenet_stats)
ar_tfm = lambda x: ((x[0], x[1]), x[1]) # auto-regression only for dtd
data.add_tfm(ar_tfm)
logging.info(f'{len(data.train_ds)} training items found.')
if not data.empty_val:
logging.info(f'{len(data.valid_ds)} valid items found.')
return data
def _get_model(config):
import importlib
names = config.model_name.split('.')
module_name, class_name = '.'.join(names[:-1]), names[-1]
cls = getattr(importlib.import_module(module_name), class_name)
model = cls(config)
logging.info(model)
return model
def _get_learner(config, data, model, local_rank=None):
strict = ifnone(config.model_strict, True)
if config.global_stage == 'pretrain-language':
metrics = [TopKTextAccuracy(
k=ifnone(config.model_k, 5),
charset_path=config.dataset_charset_path,
max_length=config.dataset_max_length + 1,
case_sensitive=config.dataset_eval_case_sensisitves,
model_eval=config.model_eval)]
else:
metrics = [TextAccuracy(
charset_path=config.dataset_charset_path,
max_length=config.dataset_max_length + 1,
case_sensitive=config.dataset_eval_case_sensisitves,
model_eval=config.model_eval)]
opt_type = getattr(torch.optim, config.optimizer_type)
learner = Learner(data, model, silent=True, model_dir='.',
true_wd=config.optimizer_true_wd,
wd=config.optimizer_wd,
bn_wd=config.optimizer_bn_wd,
path=config.global_workdir,
metrics=metrics,
opt_func=partial(opt_type, **config.optimizer_args or dict()),
loss_func=MultiLosses(one_hot=config.dataset_one_hot_y))
learner.split(lambda m: children(m))
if config.global_phase == 'train':
num_replicas = 1 if local_rank is None else torch.distributed.get_world_size()
phases = _get_training_phases(config, len(learner.data.train_dl)//num_replicas)
learner.callback_fns += [
partial(GeneralScheduler, phases=phases),
partial(GradientClipping, clip=config.optimizer_clip_grad),
partial(IterationCallback, name=config.global_name,
show_iters=config.training_show_iters,
eval_iters=config.training_eval_iters,
save_iters=config.training_save_iters,
start_iters=config.training_start_iters,
stats_iters=config.training_stats_iters)]
else:
learner.callbacks += [
DumpPrediction(learn=learner,
dataset='-'.join([Path(p).name for p in config.dataset_test_roots]),charset_path=config.dataset_charset_path,
model_eval=config.model_eval,
debug=config.global_debug,
image_only=config.global_image_only)]
learner.rank = local_rank
if local_rank is not None:
logging.info(f'Set model to distributed with rank {local_rank}.')
learner.model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(learner.model)
learner.model.to(local_rank)
learner = learner.to_distributed(local_rank)
if torch.cuda.device_count() > 1 and local_rank is None:
logging.info(f'Use {torch.cuda.device_count()} GPUs.')
learner.model = MyDataParallel(learner.model)
if config.model_checkpoint:
if Path(config.model_checkpoint).exists():
with open(config.model_checkpoint, 'rb') as f:
buffer = io.BytesIO(f.read())
learner.load(buffer, strict=strict)
else:
from distutils.dir_util import copy_tree
src = Path('/data/fangsc/model')/config.global_name
trg = Path('/output')/config.global_name
if src.exists(): copy_tree(str(src), str(trg))
learner.load(config.model_checkpoint, strict=strict)
logging.info(f'Read model from {config.model_checkpoint}')
elif config.global_phase == 'test':
learner.load(f'best-{config.global_name}', strict=strict)
logging.info(f'Read model from best-{config.global_name}')
if learner.opt_func.func.__name__ == 'Adadelta': # fastai bug, fix after 1.0.60
learner.fit(epochs=0, lr=config.optimizer_lr)
learner.opt.mom = 0.
return learner
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config', type=str, required=True,
help='path to config file')
parser.add_argument('--phase', type=str, default=None, choices=['train', 'test'])
parser.add_argument('--name', type=str, default=None)
parser.add_argument('--checkpoint', type=str, default=None)
parser.add_argument('--test_root', type=str, default=None)
parser.add_argument("--local_rank", type=int, default=None)
parser.add_argument('--debug', action='store_true', default=None)
parser.add_argument('--image_only', action='store_true', default=None)
parser.add_argument('--model_strict', action='store_false', default=None)
parser.add_argument('--model_eval', type=str, default=None,
choices=['alignment', 'vision', 'language'])
args = parser.parse_args()
config = Config(args.config)
if args.name is not None: config.global_name = args.name
if args.phase is not None: config.global_phase = args.phase
if args.test_root is not None: config.dataset_test_roots = [args.test_root]
if args.checkpoint is not None: config.model_checkpoint = args.checkpoint
if args.debug is not None: config.global_debug = args.debug
if args.image_only is not None: config.global_image_only = args.image_only
if args.model_eval is not None: config.model_eval = args.model_eval
if args.model_strict is not None: config.model_strict = args.model_strict
Logger.init(config.global_workdir, config.global_name, config.global_phase)
Logger.enable_file()
_set_random_seed(config.global_seed)
logging.info(config)
if args.local_rank is not None:
logging.info(f'Init distribution training at device {args.local_rank}.')
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
logging.info('Construct dataset.')
if config.global_stage == 'pretrain-language': data = _get_language_databaunch(config)
else: data = _get_databaunch(config)
logging.info('Construct model.')
model = _get_model(config)
logging.info('Construct learner.')
learner = _get_learner(config, data, model, args.local_rank)
if config.global_phase == 'train':
logging.info('Start training.')
learner.fit(epochs=config.training_epochs,
lr=config.optimizer_lr)
else:
logging.info('Start validate')
last_metrics = learner.validate()
log_str = f'eval loss = {last_metrics[0]:6.3f}, ' \
f'ccr = {last_metrics[1]:6.3f}, cwr = {last_metrics[2]:6.3f}, ' \
f'ted = {last_metrics[3]:6.3f}, ned = {last_metrics[4]:6.0f}, ' \
f'ted/w = {last_metrics[5]:6.3f}, '
logging.info(log_str)
if __name__ == '__main__':
main()