def main(args): # parse the options file print(f'Parsing file {args.options_file}...') opt = option.parse(args.options_file, is_train=False) opt = option.dict_to_nonedict(opt) print('Loading test images...') 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) test_loaders.append(test_loader) model = create_model(opt) generator = model.netG.module for test_loader in test_loaders: test_set_name = test_loader.dataset.opt['name'] print(f'Testing on dataset {test_set_name}') psnr_vals, ssim_vals = evaluate_generator(generator, test_loader, opt) psnr_val = np.mean(psnr_vals) ssim_val = np.mean(ssim_vals) print( f'Mean PSNR and SSIM for {test_set_name} on original model are: [{psnr_val}, {ssim_val}]' ) # The input shape is chosen arbitrarily to generate dummy input for creating quantsim object input_shapes = (1, 3, 24, 24) sim = quantsim.QuantizationSimModel( generator, input_shapes=input_shapes, quant_scheme=args.quant_scheme, default_output_bw=args.default_output_bw, default_param_bw=args.default_param_bw) evaluate_func = partial(evaluate_generator, options=opt) sim.compute_encodings(evaluate_func, test_loaders[0]) for test_loader in test_loaders: test_set_name = test_loader.dataset.opt['name'] print(f'Testing on dataset {test_set_name}') psnr_vals, ssim_vals = evaluate_generator(sim.model, test_loader, opt, output_dir=args.output_dir) psnr_val = np.mean(psnr_vals) ssim_val = np.mean(ssim_vals) print( f'Mean PSNR and SSIM for {test_set_name} on quantized model are: [{psnr_val}, {ssim_val}]' )
def main(): # parse options, set distributed setting, set ramdom seed opt = parse_options(is_train=False) torch.backends.cudnn.benchmark = True # torch.backends.cudnn.deterministic = True # mkdir and initialize loggers 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='codes', log_level=logging.INFO, log_file=log_file) logger.info(dict2str(opt)) # 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=opt['manual_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) for test_loader in test_loaders: test_set_name = test_loader.dataset.opt['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'])
def main(): # options parser = argparse.ArgumentParser() parser.add_argument('-opt', type=str, default='/home/daniel/storage/BasicSR/codes/options/train/train_cityscapeSFT.json', help='Path to option JSON file.') opt = option.parse(parser.parse_args().opt, is_train=True) opt = option.dict_to_nonedict(opt) # Convert to NoneDict, which return None for missing key. print(opt) # TODO: check if this folder setting code works if opt['path']['resume_state']: # resuming training resume_state = torch.load(opt['path']['resume_state']) else: # training from scratch resume_state = None util.mkdir_and_rename(opt['path']['experiments_root']) # rename old folder if exists util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root' and 'pretrain_model' not in key and 'resume' not in key)) # TODO: deal with the use of "logging" in this crackden writer = SummaryWriter() writer.add_text('Options', str(opt)) torch.backends.cudnn.benchmark = True # create train and val dataloader train_loader = cdl.build_custom_dataloader(opt, 'train') val_loader = cdl.build_custom_dataloader(opt, 'val') model = create_model(opt) # training current_step = 0 start_epoch = 0 total_epochs = 2000 epoch_len = len(train_loader) print('starting training') for epoch in range(start_epoch, total_epochs): tic = time.time() for _, train_data in enumerate(train_loader): current_step += 1 # model.update_learning_rate() model.feed_data(train_data) model.optimize_parameters(current_step) if current_step % 30 == 1: losses = model.get_current_losses() epoch_step = current_step % epoch_len visuals = model.get_current_visuals() gt_img = util.tensor2img(visuals['HR']) # uint8 sr_img = util.tensor2img(visuals['SR']) # uint8 writer.add_image('GT_train', gt_img, current_step, dataformats='HWC') writer.add_image('SR_train', sr_img, current_step, dataformats='HWC') print(f'step {epoch_step}/{epoch_len} || loss: {losses}') writer.add_scalars('Losses', losses, current_step) if current_step % opt['train']['val_freq'] == 0: print('validating model') avg_psnr = validate(current_step, model, opt, val_loader) writer.add_scalar('Validation loss', avg_psnr, current_step) if current_step % opt['logger']['save_checkpoint_freq'] == 0: print('Saving models and training states.') model.save(current_step) model.save_training_state(epoch, current_step) # more logging if epoch % 10 == 0: print(f'epoch {epoch} done, took {time.time()-tic}s')
def main(): # parse options, set distributed setting, set ramdom seed opt = parse_options(is_train=True) torch.backends.cudnn.benchmark = True # torch.backends.cudnn.deterministic = True # load resume states if necessary 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 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: mkdir_and_rename(osp.join('tb_logger', opt['name'])) # initialize loggers logger, tb_logger = init_loggers(opt) # create train and validation dataloaders result = create_train_val_dataloader(opt, logger) train_loader, train_sampler, val_loader, total_epochs, total_iters = result # create model if resume_state: # resume training check_resume(opt, resume_state['iter']) model = create_model(opt) 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: model = create_model(opt) 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()
# print to file and std_out simultaneously sys.stdout = PrintLogger(opt['path']['log']) print('\n**********' + util.get_timestamp() + '**********') # 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) print('Number of test images in [{:s}]: {:d}'.format( dataset_opt['name'], len(test_set))) test_loaders.append(test_loader) # Create model model = create_model(opt) for test_loader in test_loaders: test_set_name = test_loader.dataset.opt['name'] print('\nTesting [{:s}]...'.format(test_set_name)) test_start_time = time.time() dataset_dir = os.path.join(opt['path']['results_root'], test_set_name) util.mkdir(dataset_dir) test_results = OrderedDict() test_results['psnr'] = [] test_results['ssim'] = [] test_results['psnr_y'] = [] test_results['ssim_y'] = [] need_ground_truth = True