def parse_options(is_train=True):
parser = argparse.ArgumentParser()
parser.add_argument('-opt',
type=str,
required=True,
help='Path to option YAML file.')
parser.add_argument('--launcher',
choices=['none', 'pytorch', 'slurm'],
default='none',
help='job launcher')
parser.add_argument('--local_rank', type=int, default=0)
args = parser.parse_args()
opt = parse(args.opt, is_train=is_train)
# distributed settings
if args.launcher == 'none':
opt['dist'] = False
print('Disable distributed.', flush=True)
else:
opt['dist'] = True
if args.launcher == 'slurm' and 'dist_params' in opt:
init_dist(args.launcher, **opt['dist_params'])
else:
init_dist(args.launcher)
opt['rank'], opt['world_size'] = get_dist_info()
# random seed
seed = opt.get('manual_seed')
if seed is None:
seed = random.randint(1, 10000)
opt['manual_seed'] = seed
set_random_seed(seed + opt['rank'])
return opt
def main():
# options
parser = argparse.ArgumentParser()
parser.add_argument('-opt',
type=str,
required=True,
help='Path to option YAML file.')
parser.add_argument('--launcher',
choices=['none', 'pytorch', 'slurm'],
default='none',
help='job launcher')
parser.add_argument('--local_rank', type=int, default=0)
args = parser.parse_args()
opt = parse(args.opt, is_train=True)
# distributed training settings
if args.launcher == 'none': # non-distributed training
opt['dist'] = False
print('Disable distributed training.', flush=True)
else:
opt['dist'] = True
if args.launcher == 'slurm' and 'dist_params' in opt:
init_dist(args.launcher, **opt['dist_params'])
else:
init_dist(args.launcher)
rank, world_size = get_dist_info()
opt['rank'] = rank
opt['world_size'] = world_size
# load resume states if exists
if opt['path'].get('resume_state'):
device_id = torch.cuda.current_device()
resume_state = torch.load(
opt['path']['resume_state'],
map_location=lambda storage, loc: storage.cuda(device_id))
else:
resume_state = None
# mkdir and loggers
if resume_state is None:
make_exp_dirs(opt)
log_file = osp.join(opt['path']['log'],
f"train_{opt['name']}_{get_time_str()}.log")
logger = get_root_logger(logger_name='basicsr',
log_level=logging.INFO,
log_file=log_file)
logger.info(get_env_info())
logger.info(dict2str(opt))
# initialize tensorboard logger and wandb logger
tb_logger = None
if opt['logger'].get('use_tb_logger') and 'debug' not in opt['name']:
log_dir = './tb_logger/' + opt['name']
if resume_state is None and opt['rank'] == 0:
mkdir_and_rename(log_dir)
tb_logger = init_tb_logger(log_dir=log_dir)
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)
# random seed
seed = opt['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
opt['manual_seed'] = seed
logger.info(f'Random seed: {seed}')
set_random_seed(seed + rank)
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
# create train and val dataloaders
train_loader, val_loader = None, None
for phase, dataset_opt in opt['datasets'].items():
if phase == 'train':
dataset_enlarge_ratio = dataset_opt.get('dataset_enlarge_ratio', 1)
train_set = create_dataset(dataset_opt)
train_sampler = EnlargedSampler(train_set, world_size, rank,
dataset_enlarge_ratio)
train_loader = create_dataloader(train_set,
dataset_opt,
num_gpu=opt['num_gpu'],
dist=opt['dist'],
sampler=train_sampler,
seed=seed)
num_iter_per_epoch = math.ceil(
len(train_set) * dataset_enlarge_ratio /
(dataset_opt['batch_size_per_gpu'] * opt['world_size']))
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 == 'val':
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set,
dataset_opt,
num_gpu=opt['num_gpu'],
dist=opt['dist'],
sampler=None,
seed=seed)
logger.info(
f'Number of val images/folders in {dataset_opt["name"]}: '
f'{len(val_set)}')
else:
raise ValueError(f'Dataset phase {phase} is not recognized.')
assert train_loader is not None
# create model
if resume_state:
check_resume(opt, resume_state['iter']) # modify pretrain_model paths
model = create_model(opt)
# resume training
if resume_state:
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']
model.resume_training(resume_state) # handle optimizers and schedulers
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_time, iter_time = time.time(), time.time()
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_time = time.time() - data_time
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_time = time.time() - iter_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_time, 'data_time': data_time})
log_vars.update(model.get_current_log())
msg_logger(log_vars)
# save models and training states
if current_iter % opt['logger']['save_checkpoint_freq'] == 0:
logger.info('Saving models and training states.')
model.save(epoch, current_iter)
# validation
if opt.get('val') is not None and (current_iter %
opt['val']['val_freq'] == 0):
model.validation(val_loader, current_iter, tb_logger,
opt['val']['save_img'])
data_time = time.time()
iter_time = time.time()
train_data = prefetcher.next()
# end of iter
# 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:
model.validation(val_loader, current_iter, tb_logger,
opt['val']['save_img'])
if tb_logger:
tb_logger.close()
def parse_options(root_path, is_train=True):
parser = argparse.ArgumentParser()
parser.add_argument('-opt',
type=str,
required=True,
help='Path to option YAML file.')
parser.add_argument('--launcher',
choices=['none', 'pytorch', 'slurm'],
default='none',
help='job launcher')
parser.add_argument('--auto_resume', action='store_true')
parser.add_argument('--debug', action='store_true')
parser.add_argument('--local_rank', type=int, default=0)
parser.add_argument(
'--force_yml',
nargs='+',
default=None,
help='Force to update yml files. Examples: train:ema_decay=0.999')
args = parser.parse_args()
# parse yml to dict
with open(args.opt, mode='r') as f:
opt = yaml.load(f, Loader=ordered_yaml()[0])
# distributed settings
if args.launcher == 'none':
opt['dist'] = False
print('Disable distributed.', flush=True)
else:
opt['dist'] = True
if args.launcher == 'slurm' and 'dist_params' in opt:
init_dist(args.launcher, **opt['dist_params'])
else:
init_dist(args.launcher)
opt['rank'], opt['world_size'] = get_dist_info()
# random seed
seed = opt.get('manual_seed')
if seed is None:
seed = random.randint(1, 10000)
opt['manual_seed'] = seed
set_random_seed(seed + opt['rank'])
# force to update yml options
if args.force_yml is not None:
for entry in args.force_yml:
# now do not support creating new keys
keys, value = entry.split('=')
keys, value = keys.strip(), value.strip()
value = _postprocess_yml_value(value)
eval_str = 'opt'
for key in keys.split(':'):
eval_str += f'["{key}"]'
eval_str += '=value'
# using exec function
exec(eval_str)
opt['auto_resume'] = args.auto_resume
opt['is_train'] = is_train
# debug setting
if args.debug and not opt['name'].startswith('debug'):
opt['name'] = 'debug_' + opt['name']
if opt['num_gpu'] == 'auto':
opt['num_gpu'] = torch.cuda.device_count()
# datasets
for phase, dataset in opt['datasets'].items():
# for multiple datasets, e.g., val_1, val_2; test_1, test_2
phase = phase.split('_')[0]
dataset['phase'] = phase
if 'scale' in opt:
dataset['scale'] = opt['scale']
if dataset.get('dataroot_gt') is not None:
dataset['dataroot_gt'] = osp.expanduser(dataset['dataroot_gt'])
if dataset.get('dataroot_lq') is not None:
dataset['dataroot_lq'] = osp.expanduser(dataset['dataroot_lq'])
# paths
for key, val in opt['path'].items():
if (val is not None) and ('resume_state' in key
or 'pretrain_network' in key):
opt['path'][key] = osp.expanduser(val)
if is_train:
experiments_root = osp.join(root_path, 'experiments', opt['name'])
opt['path']['experiments_root'] = experiments_root
opt['path']['models'] = osp.join(experiments_root, 'models')
opt['path']['training_states'] = osp.join(experiments_root,
'training_states')
opt['path']['log'] = experiments_root
opt['path']['visualization'] = osp.join(experiments_root,
'visualization')
# change some options for debug mode
if 'debug' in opt['name']:
if 'val' in opt:
opt['val']['val_freq'] = 8
opt['logger']['print_freq'] = 1
opt['logger']['save_checkpoint_freq'] = 8
else: # test
results_root = osp.join(root_path, 'results', opt['name'])
opt['path']['results_root'] = results_root
opt['path']['log'] = results_root
opt['path']['visualization'] = osp.join(results_root, 'visualization')
return opt, args
def main():
# options
parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, help='Path to option YAML file.')
parser.add_argument('--launcher',
choices=['none', 'pytorch', 'slurm'],
default='none',
help='job launcher')
parser.add_argument('--local_rank', type=int, default=0)
args = parser.parse_args()
opt = parse(args.opt, is_train=True)
# distributed training settings
if args.launcher == 'none': # disabled distributed training
opt['dist'] = False
rank = -1
print('Disabled distributed training.', flush=True)
else:
opt['dist'] = True
if args.launcher == 'slurm' and 'dist_params' in opt:
init_dist(args.launcher, **opt['dist_params'])
else:
init_dist(args.launcher)
world_size = torch.distributed.get_world_size()
rank = torch.distributed.get_rank()
# load resume states if exists
if opt['path'].get('resume_state', None):
device_id = torch.cuda.current_device()
resume_state = torch.load(
opt['path']['resume_state'],
map_location=lambda storage, loc: storage.cuda(device_id))
else:
resume_state = None
# mkdir and loggers
if resume_state is None:
make_exp_dirs(opt)
log_file = osp.join(opt['path']['log'],
f"train_{opt['name']}_{get_time_str()}.log")
logger = get_root_logger(logger_name='basicsr',
log_level=logging.INFO,
log_file=log_file)
logger.info(get_env_info())
logger.info(dict2str(opt))
# initialize tensorboard logger and wandb logger
tb_logger = None
if opt['logger'].get('use_tb_logger') and 'debug' not in opt['name']:
tb_logger = init_tb_logger(log_dir='./tb_logger/' + opt['name'])
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)
# random seed
seed = opt['train']['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
logger.info(f'Random seed: {seed}')
set_random_seed(seed)
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
# create train and val dataloaders
train_loader, val_loader = None, None
for phase, dataset_opt in opt['datasets'].items():
if phase == 'train':
# dataset_ratio: enlarge the size of datasets for each epoch
dataset_enlarge_ratio = dataset_opt.get('dataset_enlarge_ratio', 1)
train_set = create_dataset(dataset_opt)
train_size = int(
math.ceil(len(train_set) / dataset_opt['batch_size']))
total_iters = int(opt['train']['niter'])
total_epochs = int(math.ceil(total_iters / train_size))
if opt['dist']:
train_sampler = DistIterSampler(train_set, world_size, rank,
dataset_enlarge_ratio)
total_epochs = total_iters / (train_size *
dataset_enlarge_ratio)
total_epochs = int(math.ceil(total_epochs))
else:
train_sampler = None
train_loader = create_dataloader(train_set, dataset_opt, opt,
train_sampler)
logger.info(f'Number of train images: {len(train_set)}, '
f'iters: {train_size}')
logger.info(
f'Total epochs needed: {total_epochs} for iters {total_iters}')
elif phase == 'val':
val_set = create_dataset(dataset_opt)
val_loader = create_dataloader(val_set, dataset_opt, opt, None)
logger.info(
f"Number of val images/folders in {dataset_opt['name']}: "
f'{len(val_set)}')
else:
raise NotImplementedError(f'Phase {phase} is not recognized.')
assert train_loader is not None
# create model
if resume_state:
check_resume(opt, resume_state['iter']) # modify pretrain_model paths
model = create_model(opt)
# resume training
if resume_state:
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']
model.resume_training(resume_state) # handle optimizers and schedulers
else:
current_iter = 0
start_epoch = 0
# create message logger (formatted outputs)
msg_logger = MessageLogger(opt, current_iter, tb_logger)
# training
logger.info(
f'Start training from epoch: {start_epoch}, iter: {current_iter}')
data_time, iter_time = 0, 0
for epoch in range(start_epoch, total_epochs + 1):
if opt['dist']:
train_sampler.set_epoch(epoch)
for _, train_data in enumerate(train_loader):
data_time = time.time() - data_time
current_iter += 1
if current_iter > total_iters:
break
# update learning rate
model.update_learning_rate(current_iter,
warmup_iter=opt['train']['warmup_iter'])
# training
model.feed_data(train_data)
model.optimize_parameters(current_iter)
iter_time = time.time() - iter_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_time, 'data_time': data_time})
log_vars.update(model.get_current_log())
msg_logger(log_vars)
# save models and training states
if current_iter % opt['logger']['save_checkpoint_freq'] == 0:
logger.info('Saving models and training states.')
model.save(epoch, current_iter)
# validation
if opt['val']['val_freq'] is not None and current_iter % opt[
'val']['val_freq'] == 0:
model.validation(val_loader, current_iter, tb_logger,
opt['val']['save_img'])
data_time = time.time()
iter_time = time.time()
# end of iter
# end of epoch
logger.info('End of training.')
logger.info('Saving the latest model.')
model.save(epoch=-1, current_iter=-1) # -1 for the latest
# last validation
if opt['val']['val_freq'] is not None:
model.validation(val_loader, current_iter, tb_logger,
opt['val']['save_img'])
if tb_logger:
tb_logger.close()
def main():
# options
parser = argparse.ArgumentParser()
parser.add_argument('-opt',
type=str,
required=True,
help='Path to option YAML file.')
parser.add_argument('--launcher',
choices=['none', 'pytorch', 'slurm'],
default='none',
help='job launcher')
parser.add_argument('--local_rank', type=int, default=0)
args = parser.parse_args()
opt = parse(args.opt, is_train=False)
# distributed testing settings
if args.launcher == 'none': # non-distributed testing
opt['dist'] = False
print('Disable distributed testing.', flush=True)
else:
opt['dist'] = True
if args.launcher == 'slurm' and 'dist_params' in opt:
init_dist(args.launcher, **opt['dist_params'])
else:
init_dist(args.launcher)
rank, world_size = get_dist_info()
opt['rank'] = rank
opt['world_size'] = world_size
make_exp_dirs(opt)
log_file = osp.join(opt['path']['log'],
f"test_{opt['name']}_{get_time_str()}.log")
logger = get_root_logger(logger_name='basicsr',
log_level=logging.INFO,
log_file=log_file)
logger.info(get_env_info())
logger.info(dict2str(opt))
# random seed
seed = opt['manual_seed']
if seed is None:
seed = random.randint(1, 10000)
opt['manual_seed'] = seed
logger.info(f'Random seed: {seed}')
set_random_seed(seed + rank)
torch.backends.cudnn.benchmark = True
# torch.backends.cudnn.deterministic = True
# create test dataset and dataloader
test_loaders = []
for phase, dataset_opt in sorted(opt['datasets'].items()):
test_set = create_dataset(dataset_opt)
test_loader = create_dataloader(test_set,
dataset_opt,
num_gpu=opt['num_gpu'],
dist=opt['dist'],
sampler=None,
seed=seed)
logger.info(
f"Number of test images in {dataset_opt['name']}: {len(test_set)}")
test_loaders.append(test_loader)
# create model
model = create_model(opt)
dummy_input = torch.zeros((1, 3, 800, 800)).to(model.device)
dummy_psf = torch.zeros((1, 5, 1, 1)).to(model.device)
summary(model.net_g, dummy_input, dummy_psf)
for test_loader in test_loaders:
test_set_name = test_loader.dataset.opt['name']
mkdir_or_exist(osp.join(opt['path']['visualization'], test_set_name))
logger.info(f'Testing {test_set_name}...')
model.validation(test_loader,
current_iter=opt['name'],
tb_logger=None,
save_img=opt['val']['save_img'])
parser.add_argument('--truncation_mean', type=int, default=4096)
parser.add_argument(
'--ckpt',
type=str,
default= # noqa: E251
'experiments/pretrained_models/StyleGAN/stylegan2_ffhq_config_f_1024_official-b09c3668.pth' # noqa: E501
)
parser.add_argument('--channel_multiplier', type=int, default=2)
parser.add_argument('--randomize_noise', type=bool, default=True)
args = parser.parse_args()
args.latent = 512
args.n_mlp = 8
os.makedirs('samples', exist_ok=True)
set_random_seed(2020)
g_ema = StyleGAN2Generator(
args.size,
args.latent,
args.n_mlp,
channel_multiplier=args.channel_multiplier).to(device)
checkpoint = torch.load(args.ckpt)['params_ema']
g_ema.load_state_dict(checkpoint)
if args.truncation < 1:
with torch.no_grad():
mean_latent = g_ema.mean_latent(args.truncation_mean)
else:
mean_latent = None
