def test_glean(): model_cfg = dict(type='GLEAN', generator=dict(type='GLEANStyleGANv2', in_size=16, out_size=64, style_channels=512), discriminator=dict(type='StyleGAN2Discriminator', in_size=64), pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean'), gan_loss=dict(type='GANLoss', gan_type='vanilla', real_label_val=1.0, fake_label_val=0, loss_weight=5e-3)) train_cfg = None test_cfg = mmcv.Config(dict(metrics=['PSNR'], crop_border=0)) # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # prepare data inputs = torch.rand(1, 3, 16, 16) targets = torch.rand(1, 3, 64, 64) data_batch = {'lq': inputs, 'gt': targets} restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) meta = [{'lq_path': ''}] # test forward_test (cpu) with pytest.raises(ValueError): # iteration is not None or number with torch.no_grad(): restorer(**data_batch, test_mode=True, save_image=True, meta=meta, iteration='1') with pytest.raises(AssertionError): # test with metric but gt is None with torch.no_grad(): data_batch.pop('gt') restorer(**data_batch, test_mode=True) # test forward_test (gpu) if torch.cuda.is_available(): data_batch = {'lq': inputs.cuda(), 'gt': targets.cuda()} restorer = restorer.cuda() with pytest.raises(ValueError): # iteration is not None or number with torch.no_grad(): restorer(**data_batch, test_mode=True, save_image=True, meta=meta, iteration='1') with pytest.raises(AssertionError): # test with metric but gt is None with torch.no_grad(): data_batch.pop('gt') restorer(**data_batch, test_mode=True)
def test_indexnet(): model_cfg, _, test_cfg = _get_model_cfg( 'indexnet/indexnet_mobv2_1x16_78k_comp1k.py') model_cfg['pretrained'] = None # test indexnet inference with torch.no_grad(): indexnet = build_model(model_cfg, train_cfg=None, test_cfg=test_cfg) indexnet.eval() input_test = _demo_input_test((64, 64)) output_test = indexnet(**input_test, test_mode=True) assert isinstance(output_test['pred_alpha'], np.ndarray) assert output_test['pred_alpha'].shape == (64, 64) assert_dict_keys_equal(output_test['eval_result'], ['SAD', 'MSE', 'GRAD', 'CONN']) # test inference with gpu if torch.cuda.is_available(): indexnet = build_model( model_cfg, train_cfg=None, test_cfg=test_cfg).cuda() indexnet.eval() input_test = _demo_input_test((64, 64), cuda=True) output_test = indexnet(**input_test, test_mode=True) assert isinstance(output_test['pred_alpha'], np.ndarray) assert output_test['pred_alpha'].shape == (64, 64) assert_dict_keys_equal(output_test['eval_result'], ['SAD', 'MSE', 'GRAD', 'CONN']) # test forward train though we do not guarantee the training for present model_cfg.loss_alpha = None model_cfg.loss_comp = dict(type='L1CompositionLoss') indexnet = build_model( model_cfg, train_cfg=mmcv.ConfigDict(train_backbone=True), test_cfg=test_cfg) input_train = _demo_input_train((64, 64), batch_size=2) output_train = indexnet(**input_train) assert output_train['num_samples'] == 2 assert_dict_keys_equal(output_train['losses'], ['loss_comp']) if torch.cuda.is_available(): model_cfg.loss_alpha = dict(type='L1Loss') model_cfg.loss_comp = None indexnet = build_model( model_cfg, train_cfg=mmcv.ConfigDict(train_backbone=True), test_cfg=test_cfg).cuda() input_train = _demo_input_train((64, 64), batch_size=2, cuda=True) output_train = indexnet(**input_train) assert output_train['num_samples'] == 2 assert_dict_keys_equal(output_train['losses'], ['loss_alpha']) # test forward_dummy indexnet.cpu().eval() inputs = torch.ones((1, 4, 32, 32)) indexnet.forward_dummy(inputs)
def test_gl_inpaintor(): cfg = Config.fromfile( Path(__file__).parent.joinpath('data/inpaintor_config/gl_test.py')) gl = build_model(cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg) assert gl.__class__.__name__ == 'GLInpaintor' if torch.cuda.is_available(): gt_img = torch.randn(1, 3, 256, 256) mask = torch.zeros_like(gt_img)[:, 0:1, ...] mask[..., 100:210, 100:210] = 1. masked_img = gt_img * (1. - mask) mask_bbox = torch.tensor([[100, 100, 110, 110]]) gl.cuda() data_batch = dict( gt_img=gt_img.cuda(), mask=mask.cuda(), masked_img=masked_img.cuda(), mask_bbox=mask_bbox.cuda()) optim_g = torch.optim.SGD(gl.generator.parameters(), lr=0.1) optim_d = torch.optim.SGD(gl.disc.parameters(), lr=0.1) optim_dict = dict(generator=optim_g, disc=optim_d) for i in range(5): outputs = gl.train_step(data_batch, optim_dict) if i <= 2: assert 'loss_l1_hole' in outputs['log_vars'] assert 'fake_loss' not in outputs['log_vars'] assert 'real_loss' not in outputs['log_vars'] assert 'loss_g_fake' not in outputs['log_vars'] elif i == 3: assert 'loss_l1_hole' not in outputs['log_vars'] assert 'fake_loss' in outputs['log_vars'] assert 'real_loss' in outputs['log_vars'] assert 'loss_g_fake' not in outputs['log_vars'] else: assert 'loss_l1_hole' in outputs['log_vars'] assert 'fake_loss' in outputs['log_vars'] assert 'real_loss' in outputs['log_vars'] assert 'loss_g_fake' in outputs['log_vars'] gl_dirty = build_model( cfg.model_dirty, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg) gl_dirty.cuda() res, loss = gl_dirty.generator_loss(gt_img, gt_img, gt_img, data_batch) assert len(loss) == 0
def __init__(self, trt_file: str, cfg: Any, device_id: int): super().__init__() base_model = build_model( cfg.model, train_cfg=None, test_cfg=cfg.test_cfg) if isinstance(base_model, BasicRestorer): WraperClass = TensorRTRestorer self.wraper = WraperClass(base_model, trt_file, device_id)
def main(): args = parse_args() if len(args.shape) == 1: input_shape = (3, args.shape[0], args.shape[0]) elif len(args.shape) == 2: input_shape = (3, ) + tuple(args.shape) elif len(args.shape) == 3: input_shape = tuple(args.shape) else: raise ValueError('invalid input shape') cfg = Config.fromfile(args.config) model = build_model(cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg).cuda() model.eval() if hasattr(model, 'forward_dummy'): model.forward = model.forward_dummy else: raise NotImplementedError( 'FLOPs counter is currently not currently supported ' f'with {model.__class__.__name__}') flops, params = get_model_complexity_info(model, input_shape) split_line = '=' * 30 print(f'{split_line}\nInput shape: {input_shape}\n' f'Flops: {flops}\nParams: {params}\n{split_line}') print('!!!Please be cautious if you use the results in papers. ' 'You may need to check if all ops are supported and verify that the ' 'flops computation is correct.')
def init_model(config, checkpoint=None, device='cuda:0'): """Initialize a model from config file. Args: config (str or :obj:`mmcv.Config`): Config file path or the config object. checkpoint (str, optional): Checkpoint path. If left as None, the model will not load any weights. device (str): Which device the model will deploy. Default: 'cuda:0'. Returns: nn.Module: The constructed model. """ if isinstance(config, str): config = mmcv.Config.fromfile(config) elif not isinstance(config, mmcv.Config): raise TypeError('config must be a filename or Config object, ' f'but got {type(config)}') config.model.pretrained = None config.test_cfg.metrics = None model = build_model(config.model, test_cfg=config.test_cfg) if checkpoint is not None: checkpoint = load_checkpoint(model, checkpoint) model.cfg = config # save the config in the model for convenience model.to(device) model.eval() return model
def test_gca(): model_cfg, train_cfg, test_cfg = _get_model_cfg( 'gca/gca_r34_4x10_200k_comp1k.py') model_cfg['pretrained'] = None # test model forward in train mode model = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) inputs = _demo_input_train((64, 64), batch_size=2) inputs['trimap'] = inputs['trimap'].expand_as(inputs['merged']) inputs['meta'][0]['to_onehot'] = True outputs = model(inputs['merged'], inputs['trimap'], inputs['meta'], inputs['alpha']) assert outputs['num_samples'] == 2 assert_dict_keys_equal(outputs['losses'], ['loss']) if torch.cuda.is_available(): model = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) model.cuda() inputs = _demo_input_train((64, 64), batch_size=2, cuda=True) inputs['trimap'] = inputs['trimap'].expand_as(inputs['merged']) inputs['meta'][0]['to_onehot'] = True outputs = model(inputs['merged'], inputs['trimap'], inputs['meta'], inputs['alpha']) assert outputs['num_samples'] == 2 assert_dict_keys_equal(outputs['losses'], ['loss']) # test model forward in test mode with torch.no_grad(): model_cfg.backbone.encoder.in_channels = 4 model = build_model(model_cfg, train_cfg=None, test_cfg=test_cfg) inputs = _demo_input_test((64, 64)) outputs = model(**inputs, test_mode=True) assert_dict_keys_equal(outputs['eval_result'], ['SAD', 'MSE', 'GRAD', 'CONN']) if torch.cuda.is_available(): model = build_model(model_cfg, train_cfg=None, test_cfg=test_cfg) model.cuda() inputs = _demo_input_test((64, 64), cuda=True) outputs = model(**inputs, test_mode=True) assert_dict_keys_equal(outputs['eval_result'], ['SAD', 'MSE', 'GRAD', 'CONN']) # test forward_dummy model.cpu().eval() inputs = torch.ones((1, 4, 32, 32)) model.forward_dummy(inputs)
def check_params(cfg): model = build_model(cfg.model).cuda() # print(model) num_param = sum(p.numel() for p in model.parameters()) if num_param > 1821085: raise Exception('model parameters exceed limit') else: print( 'there are total of {} parameters in the model'.format(num_param))
def main(): print('settings:\n', args) #annotate training data for 1st time if args.annotate: train_annotation() #change config config_path = 'configs/restorers/srresnet_srgan/msrresnet_x4c64b16_g1_1000k_div2k.py' cfg = change_config(config_path) check_params(cfg) # Initialize distributed training (only need to initialize once), comment it if have already run this part os.environ['RANK'] = '0' os.environ['WORLD_SIZE'] = '1' os.environ['MASTER_ADDR'] = '127.0.0.1' os.environ['MASTER_PORT'] = '29500' #'50297' init_dist('pytorch', **cfg.dist_params) # Build dataset datasets = [build_dataset(cfg.data.train)] # Build the SRCNN model model = build_model(cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg) # Create work_dir mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir)) # Meta information meta = dict() # if cfg.get('exp_name', None) is None: # cfg['exp_name'] = osp.splitext(osp.basename(cfg.work_dir))[0] meta['exp_name'] = '_'.join([ 'bs' + str(args.bs), 'iter' + str(args.iter), 'block' + str(args.num_blocks), args.loss ]) meta['mmedit Version'] = mmedit.__version__ meta['seed'] = 0 meta['start_time'] = datetime.now().strftime("%d/%m/%Y %H:%M:%S") # Train the model train_model(model, datasets, cfg, distributed=True, validate=True, meta=meta)
def __init__(self, onnx_file, cfg, device_id): super(ONNXRuntimeEditing, self).__init__() ort_custom_op_path = '' try: from mmcv.ops import get_onnxruntime_op_path ort_custom_op_path = get_onnxruntime_op_path() except (ImportError, ModuleNotFoundError): warnings.warn('If input model has custom op from mmcv, \ you may have to build mmcv with ONNXRuntime from source.') session_options = ort.SessionOptions() # register custom op for onnxruntime if osp.exists(ort_custom_op_path): session_options.register_custom_ops_library(ort_custom_op_path) sess = ort.InferenceSession(onnx_file, session_options) providers = ['CPUExecutionProvider'] options = [{}] is_cuda_available = ort.get_device() == 'GPU' if is_cuda_available: providers.insert(0, 'CUDAExecutionProvider') options.insert(0, {'device_id': device_id}) sess.set_providers(providers, options) self.sess = sess self.device_id = device_id self.io_binding = sess.io_binding() self.output_names = [_.name for _ in sess.get_outputs()] base_model = build_model(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg) if isinstance(base_model, BaseMattor): WraperClass = ONNXRuntimeMattor elif isinstance(base_model, BasicRestorer): WraperClass = ONNXRuntimeRestorer self.wraper = WraperClass(self.sess, self.io_binding, self.output_names, base_model)
def test_liif(): model_cfg = dict(type='LIIF', generator=dict(type='LIIFEDSR', encoder=dict(type='EDSR', in_channels=3, out_channels=3, mid_channels=64, num_blocks=16), imnet=dict( type='MLPRefiner', in_dim=64, out_dim=3, hidden_list=[256, 256, 256, 256]), local_ensemble=True, feat_unfold=True, cell_decode=True, eval_bsize=30000), rgb_mean=(0.4488, 0.4371, 0.4040), rgb_std=(1., 1., 1.), pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean')) scale_max = 4 train_cfg = None test_cfg = Config(dict(metrics=['PSNR', 'SSIM'], crop_border=scale_max)) # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test attributes assert restorer.__class__.__name__ == 'LIIF' # prepare data inputs = torch.rand(1, 3, 22, 11) targets = torch.rand(1, 128 * 64, 3) coord = torch.rand(1, 128 * 64, 2) cell = torch.rand(1, 128 * 64, 2) data_batch = {'lq': inputs, 'gt': targets, 'coord': coord, 'cell': cell} # prepare optimizer optim_cfg = dict(type='Adam', lr=1e-4, betas=(0.9, 0.999)) optimizer = obj_from_dict(optim_cfg, torch.optim, dict(params=restorer.parameters())) # test train_step and forward_test (cpu) outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert outputs['results']['lq'].shape == data_batch['lq'].shape assert outputs['results']['gt'].shape == data_batch['gt'].shape assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 128 * 64, 3) # test train_step and forward_test (gpu) if torch.cuda.is_available(): restorer = restorer.cuda() data_batch = { 'lq': inputs.cuda(), 'gt': targets.cuda(), 'coord': coord.cuda(), 'cell': cell.cuda() } # train_step optimizer = obj_from_dict(optim_cfg, torch.optim, dict(params=restorer.parameters())) outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert outputs['results']['lq'].shape == data_batch['lq'].shape assert outputs['results']['gt'].shape == data_batch['gt'].shape assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 128 * 64, 3) # val_step result = restorer.val_step(data_batch, meta=[{'gt_path': ''}]) assert isinstance(result, dict) assert isinstance(result['eval_result'], dict) assert result['eval_result'].keys() == set({'PSNR', 'SSIM'}) assert isinstance(result['eval_result']['PSNR'], np.float64) assert isinstance(result['eval_result']['SSIM'], np.float64)
def test_mattor_wrapper(): try: import onnxruntime as ort from mmedit.core.export.wrappers import (ONNXRuntimeEditing, ONNXRuntimeMattor) except ImportError: pytest.skip('ONNXRuntime is not available.') onnx_path = 'tmp.onnx' train_cfg = None test_cfg = dict(refine=False, metrics=['SAD', 'MSE', 'GRAD', 'CONN']) cfg = dict( model=dict( type='DIM', backbone=dict( type='SimpleEncoderDecoder', encoder=dict(type='VGG16', in_channels=4), decoder=dict(type='PlainDecoder')), pretrained='open-mmlab://mmedit/vgg16', loss_alpha=dict(type='CharbonnierLoss', loss_weight=0.5), loss_comp=dict(type='CharbonnierCompLoss', loss_weight=0.5)), train_cfg=train_cfg, test_cfg=test_cfg) cfg = mmcv.Config(cfg) pytorch_model = build_model( cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg) img_shape = (32, 32) merged = torch.rand(1, 3, img_shape[1], img_shape[0]) trimap = torch.rand(1, 1, img_shape[1], img_shape[0]) data_batch = {'merged': merged, 'trimap': trimap} inputs = torch.cat([merged, trimap], dim=1) pytorch_model.forward = pytorch_model.forward_dummy with torch.no_grad(): torch.onnx.export( pytorch_model, inputs, onnx_path, input_names=['input'], output_names=['output'], export_params=True, keep_initializers_as_inputs=False, verbose=False, opset_version=11) wrap_model = ONNXRuntimeEditing(onnx_path, cfg, 0) os.remove(onnx_path) assert isinstance(wrap_model.wrapper, ONNXRuntimeMattor) if ort.get_device() == 'GPU': merged = merged.cuda() trimap = trimap.cuda() data_batch = {'merged': merged, 'trimap': trimap} ori_alpha = np.random.random(img_shape).astype(np.float32) ori_trimap = np.random.randint(256, size=img_shape).astype(np.float32) data_batch['meta'] = [ dict( ori_alpha=ori_alpha, ori_trimap=ori_trimap, merged_ori_shape=img_shape) ] with torch.no_grad(): outputs = wrap_model(**data_batch, test_mode=True) assert isinstance(outputs, dict) assert 'pred_alpha' in outputs pred_alpha = outputs['pred_alpha'] assert isinstance(pred_alpha, np.ndarray) assert pred_alpha.shape == img_shape
def test_pconv_inpaintor(init_weights): cfg = Config.fromfile( 'tests/data/inpaintor_config/pconv_inpaintor_test.py') if torch.cuda.is_available(): pconv_inpaintor = build_model(cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg) assert pconv_inpaintor.__class__.__name__ == 'PConvInpaintor' pconv_inpaintor.cuda() gt_img = torch.randn((1, 3, 256, 256)).cuda() mask = torch.zeros_like(gt_img) mask[..., 50:160, 100:210] = 1. masked_img = gt_img * (1. - mask) data_batch = dict(gt_img=gt_img, mask=mask, masked_img=masked_img) optim_g = torch.optim.SGD(pconv_inpaintor.generator.parameters(), lr=0.1) optim_dict = dict(generator=optim_g) outputs = pconv_inpaintor.train_step(data_batch, optim_dict) assert outputs['results']['fake_res'].shape == (1, 3, 256, 256) assert outputs['results']['final_mask'].shape == (1, 3, 256, 256) assert 'loss_l1_hole' in outputs['log_vars'] assert 'loss_l1_valid' in outputs['log_vars'] assert 'loss_tv' in outputs['log_vars'] # test forward dummy res = pconv_inpaintor.forward_dummy( torch.cat([masked_img, mask], dim=1)) assert res.shape == (1, 3, 256, 256) # test forward test w/o save image outputs = pconv_inpaintor.forward_test(masked_img[0:1], mask[0:1], gt_img=gt_img[0:1, ...]) assert 'eval_result' in outputs assert outputs['eval_result']['l1'] > 0 assert outputs['eval_result']['psnr'] > 0 assert outputs['eval_result']['ssim'] > 0 # test forward test w/o eval metrics pconv_inpaintor.test_cfg = dict() pconv_inpaintor.eval_with_metrics = False outputs = pconv_inpaintor.forward_test(masked_img[0:1], mask[0:1]) for key in ['fake_res', 'fake_img']: assert outputs[key].size() == (1, 3, 256, 256) # test forward test w/ save image with tempfile.TemporaryDirectory() as tmpdir: outputs = pconv_inpaintor.forward_test( masked_img[0:1], mask[0:1], save_image=True, save_path=tmpdir, iteration=4396, meta=[dict(gt_img_path='igccc.png')]) assert os.path.exists(os.path.join(tmpdir, 'igccc_4396.png')) # test forward test w/ save image w/ gt_img with tempfile.TemporaryDirectory() as tmpdir: outputs = pconv_inpaintor.forward_test( masked_img[0:1], mask[0:1], save_image=True, save_path=tmpdir, meta=[dict(gt_img_path='igccc.png')], gt_img=gt_img[0:1, ...]) assert os.path.exists(os.path.join(tmpdir, 'igccc.png')) with pytest.raises(AssertionError): outputs = pconv_inpaintor.forward_test(masked_img[0:1], mask[0:1], save_image=True, save_path=tmpdir, iteration=4396, gt_img=gt_img[0:1, ...]) with pytest.raises(AssertionError): outputs = pconv_inpaintor.forward_test( masked_img[0:1], mask[0:1], save_image=True, save_path=None, iteration=4396, meta=[dict(gt_img_path='igccc.png')], gt_img=gt_img[0:1, ...]) # reset mock to clear some memory usage init_weights.reset_mock()
def test_cain(): model_cfg = dict( type='CAIN', generator=dict(type='CAINNet'), pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean')) train_cfg = None test_cfg = None # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test attributes assert restorer.__class__.__name__ == 'CAIN' assert isinstance(restorer.generator, CAINNet) assert isinstance(restorer.pixel_loss, L1Loss) # prepare data inputs = torch.rand(1, 2, 3, 128, 128) target = torch.rand(1, 3, 128, 128) data_batch = {'inputs': inputs, 'target': target, 'meta': [{'key': '001'}]} # prepare optimizer optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.9, 0.999)) optimizer = { 'generator': obj_from_dict(optim_cfg, torch.optim, dict(params=restorer.parameters())) } # test forward_test with torch.no_grad(): outputs = restorer.forward_test(**data_batch) assert torch.equal(outputs['inputs'], data_batch['inputs']) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 128, 128) # test train_step outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['inputs'], data_batch['inputs']) assert torch.equal(outputs['results']['target'], data_batch['target']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 128, 128) # test train_step and forward_test (gpu) if torch.cuda.is_available(): restorer = restorer.cuda() optimizer['generator'] = obj_from_dict( optim_cfg, torch.optim, dict(params=restorer.parameters())) data_batch = { 'inputs': inputs.cuda(), 'target': target.cuda(), 'meta': [{ 'key': '001' }] } # forward_test with torch.no_grad(): outputs = restorer.forward_test(**data_batch) assert torch.equal(outputs['inputs'], data_batch['inputs'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 128, 128) # train_step outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['inputs'], data_batch['inputs'].cpu()) assert torch.equal(outputs['results']['target'], data_batch['target'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 128, 128) # test with metric and save image test_cfg = dict(metrics=('PSNR', 'SSIM'), crop_border=0) test_cfg = mmcv.Config(test_cfg) data_batch = { 'inputs': inputs, 'target': target, 'meta': [{ 'key': 'fake_path/fake_name' }] } restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) with pytest.raises(AssertionError): # evaluation with metrics must have target images restorer(inputs=inputs, test_mode=True) with tempfile.TemporaryDirectory() as tmpdir: outputs = restorer( **data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=None) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) outputs = restorer( **data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=100) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float)
def test_restorer_wrapper(): try: import onnxruntime as ort from mmedit.core.export.wrappers import (ONNXRuntimeEditing, ONNXRuntimeRestorer) except ImportError: pytest.skip('ONNXRuntime is not available.') onnx_path = 'tmp.onnx' scale = 4 train_cfg = None test_cfg = None cfg = dict( model=dict( type='BasicRestorer', generator=dict( type='SRCNN', channels=(3, 4, 2, 3), kernel_sizes=(9, 1, 5), upscale_factor=scale), pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean')), train_cfg=train_cfg, test_cfg=test_cfg) cfg = mmcv.Config(cfg) pytorch_model = build_model( cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg) # prepare data inputs = torch.rand(1, 3, 2, 2) targets = torch.rand(1, 3, 8, 8) data_batch = {'lq': inputs, 'gt': targets} pytorch_model.forward = pytorch_model.forward_dummy with torch.no_grad(): torch.onnx.export( pytorch_model, inputs, onnx_path, input_names=['input'], output_names=['output'], export_params=True, keep_initializers_as_inputs=False, verbose=False, opset_version=11) wrap_model = ONNXRuntimeEditing(onnx_path, cfg, 0) # os.remove(onnx_path) assert isinstance(wrap_model.wrapper, ONNXRuntimeRestorer) if ort.get_device() == 'GPU': data_batch = {'lq': inputs.cuda(), 'gt': targets.cuda()} with torch.no_grad(): outputs = wrap_model(**data_batch, test_mode=True) assert isinstance(outputs, dict) assert 'output' in outputs output = outputs['output'] assert isinstance(output, torch.Tensor) assert output.shape == targets.shape
def main(): args = parse_args() checkpoint_list = os.listdir(args.checkpoint_dir) print(checkpoint_list) for checkpoint in checkpoint_list: if '.pth' in checkpoint: cfg = mmcv.Config.fromfile(args.config) # set cudnn_benchmark if cfg.get('cudnn_benchmark', False): torch.backends.cudnn.benchmark = True cfg.model.pretrained = None # init distributed env first, since logger depends on the dist info. if args.launcher == 'none': distributed = False else: distributed = True init_dist(args.launcher, **cfg.dist_params) rank, _ = get_dist_info() # set random seeds if args.seed is not None: if rank == 0: print('set random seed to', args.seed) set_random_seed(args.seed, deterministic=args.deterministic) # build the dataloader # TODO: support multiple images per gpu (only minor changes are needed) dataset = build_dataset(cfg.data.test) data_loader = build_dataloader(dataset, samples_per_gpu=1, workers_per_gpu=cfg.data.get( 'val_workers_per_gpu', cfg.data.workers_per_gpu), dist=distributed, shuffle=False) # build the model and load checkpoint model = build_model(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg) args.save_image = args.save_path is not None # distributed test find_unused_parameters = cfg.get('find_unused_parameters', False) model = DistributedDataParallelWrapper( model, device_ids=[torch.cuda.current_device()], broadcast_buffers=False, find_unused_parameters=find_unused_parameters) device_id = torch.cuda.current_device() _ = load_checkpoint( model, os.path.join(args.checkpoint_dir, checkpoint), map_location=lambda storage, loc: storage.cuda(device_id)) outputs = multi_gpu_test(model, data_loader, args.tmpdir, args.gpu_collect, save_path=args.save_path, save_image=args.save_image) if rank == 0: # print metrics stats = dataset.evaluate(outputs) write_file = open( os.path.join(args.checkpoint_dir, 'eval_result_new.txt'), 'a') for stat in stats: print('{}: Eval-{}: {}'.format(checkpoint, stat, stats[stat])) write_file.write('{}: Eval-{}: {} '.format( checkpoint, stat, stats[stat])) write_file.write('\n') write_file.close() # save result pickle if args.out: print('writing results to {}'.format(args.out)) mmcv.dump(outputs, args.out)
def test_basic_interpolator(): model_cfg = dict(type='BasicInterpolator', generator=dict(type='InterpolateExample'), pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean')) train_cfg = None test_cfg = None # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test attributes assert restorer.__class__.__name__ == 'BasicInterpolator' assert isinstance(restorer.generator, InterpolateExample) assert isinstance(restorer.pixel_loss, L1Loss) # prepare data inputs = torch.rand(1, 2, 3, 20, 20) target = torch.rand(1, 3, 20, 20) data_batch = {'inputs': inputs, 'target': target} # prepare optimizer optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.9, 0.999)) optimizer = { 'generator': obj_from_dict(optim_cfg, torch.optim, dict(params=restorer.parameters())) } # test forward train outputs = restorer(**data_batch, test_mode=False) assert isinstance(outputs, dict) assert isinstance(outputs['losses'], dict) assert isinstance(outputs['losses']['loss_pix'], torch.FloatTensor) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['inputs'], data_batch['inputs']) assert torch.equal(outputs['results']['target'], data_batch['target']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 20, 20) # test forward_test with torch.no_grad(): restorer.val_step(data_batch) outputs = restorer(**data_batch, test_mode=True) assert torch.equal(outputs['inputs'], data_batch['inputs']) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 20, 20) # test forward_dummy with torch.no_grad(): output = restorer.forward_dummy(data_batch['inputs']) assert torch.is_tensor(output) assert output.size() == (1, 3, 20, 20) # test train_step outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['inputs'], data_batch['inputs']) assert torch.equal(outputs['results']['target'], data_batch['target']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 20, 20) # test train_step and forward_test (gpu) if torch.cuda.is_available(): restorer = restorer.cuda() optimizer['generator'] = obj_from_dict( optim_cfg, torch.optim, dict(params=restorer.parameters())) data_batch = {'inputs': inputs.cuda(), 'target': target.cuda()} # test forward train outputs = restorer(**data_batch, test_mode=False) assert isinstance(outputs, dict) assert isinstance(outputs['losses'], dict) assert isinstance(outputs['losses']['loss_pix'], torch.cuda.FloatTensor) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['inputs'], data_batch['inputs'].cpu()) assert torch.equal(outputs['results']['target'], data_batch['target'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 20, 20) # forward_test with torch.no_grad(): restorer.val_step(data_batch) outputs = restorer(**data_batch, test_mode=True) assert torch.equal(outputs['inputs'], data_batch['inputs'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 20, 20) # train_step outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['inputs'], data_batch['inputs'].cpu()) assert torch.equal(outputs['results']['target'], data_batch['target'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 20, 20) # test with metric and save image test_cfg = dict(metrics=('PSNR', 'SSIM'), crop_border=0) test_cfg = mmcv.Config(test_cfg) data_batch = { 'inputs': inputs, 'target': target, 'meta': [{ 'key': '000001/0000', 'target_path': 'fake_path/fake_name.png' }] } restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) with pytest.raises(AssertionError): # evaluation with metrics must have target images restorer(inputs=inputs, test_mode=True) with tempfile.TemporaryDirectory() as tmpdir: outputs = restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=None) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) outputs = restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=100) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) outputs = restorer( inputs=inputs, target=target, meta=[{ 'key': '000001/0000', 'inputs_path': ['fake_path/fake_name.png', 'fake_path/fake_name.png'] }], test_mode=True, save_image=True, save_path=tmpdir, iteration=100) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) with pytest.raises(ValueError): # iteration should be number or None restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration='100') # test forward_test when output.shape==5 model_cfg = dict(type='BasicInterpolator', generator=dict(type='InterpolateExample2'), pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean')) train_cfg = None test_cfg = None restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) outputs = restorer( inputs=inputs, target=target.unsqueeze(1), meta=[{ 'key': '000001/0000', 'inputs_path': ['fake_path/fake_name.png', 'fake_path/fake_name.png'] }], test_mode=True, save_image=True, save_path=tmpdir, iteration=100) outputs = restorer( inputs=inputs, target=target.unsqueeze(1), meta=[{ 'key': '000001/0000', 'inputs_path': ['fake_path/fake_name.png', 'fake_path/fake_name.png'] }], test_mode=True, save_image=True, save_path=tmpdir, iteration=None) with pytest.raises(ValueError): # iteration should be number or None restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration='100') # test merge_frames input_tensors = torch.rand(2, 2, 3, 256, 256) output_tensors = torch.rand(2, 1, 3, 256, 256) result = restorer.merge_frames(input_tensors, output_tensors) assert isinstance(result, list) assert len(result) == 5 assert result[0].shape == (256, 256, 3) # test split_frames tensors = torch.rand(1, 10, 3, 256, 256) result = restorer.split_frames(tensors) assert isinstance(result, torch.Tensor) assert result.shape == (9, 2, 3, 256, 256) # test evaluate 5d output test_cfg = dict(metrics=('PSNR', 'SSIM'), crop_border=0) test_cfg = mmcv.Config(test_cfg) restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) output = torch.rand(1, 2, 3, 256, 256) target = torch.rand(1, 2, 3, 256, 256) restorer.evaluate(output, target)
def test_pix2pix(): model_cfg = dict(type='Pix2Pix', generator=dict(type='UnetGenerator', in_channels=3, out_channels=3, num_down=8, base_channels=64, norm_cfg=dict(type='BN'), use_dropout=True, init_cfg=dict(type='normal', gain=0.02)), discriminator=dict(type='PatchDiscriminator', in_channels=6, base_channels=64, num_conv=3, norm_cfg=dict(type='BN'), init_cfg=dict(type='normal', gain=0.02)), gan_loss=dict(type='GANLoss', gan_type='vanilla', real_label_val=1.0, fake_label_val=0, loss_weight=1.0), pixel_loss=dict(type='L1Loss', loss_weight=100.0, reduction='mean')) train_cfg = None test_cfg = None # build synthesizer synthesizer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test checking gan loss cannot be None with pytest.raises(AssertionError): bad_model_cfg = copy.deepcopy(model_cfg) bad_model_cfg['gan_loss'] = None _ = build_model(bad_model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test attributes assert synthesizer.__class__.__name__ == 'Pix2Pix' assert isinstance(synthesizer.generator, UnetGenerator) assert isinstance(synthesizer.discriminator, PatchDiscriminator) assert isinstance(synthesizer.gan_loss, GANLoss) assert isinstance(synthesizer.pixel_loss, L1Loss) assert synthesizer.train_cfg is None assert synthesizer.test_cfg is None # prepare data inputs = torch.rand(1, 3, 256, 256) targets = torch.rand(1, 3, 256, 256) data_batch = {'img_a': inputs, 'img_b': targets} img_meta = {} img_meta['img_a_path'] = 'img_a_path' img_meta['img_b_path'] = 'img_b_path' data_batch['meta'] = [img_meta] # prepare optimizer optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.5, 0.999)) optimizer = { 'generator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(synthesizer, 'generator').parameters())), 'discriminator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(synthesizer, 'discriminator').parameters())) } # test forward_dummy with torch.no_grad(): output = synthesizer.forward_dummy(data_batch['img_a']) assert torch.is_tensor(output) assert output.size() == (1, 3, 256, 256) # test forward_test with torch.no_grad(): outputs = synthesizer(inputs, targets, [img_meta], test_mode=True) assert torch.equal(outputs['real_a'], data_batch['img_a']) assert torch.equal(outputs['real_b'], data_batch['img_b']) assert torch.is_tensor(outputs['fake_b']) assert outputs['fake_b'].size() == (1, 3, 256, 256) # val_step with torch.no_grad(): outputs = synthesizer.val_step(data_batch) assert torch.equal(outputs['real_a'], data_batch['img_a']) assert torch.equal(outputs['real_b'], data_batch['img_b']) assert torch.is_tensor(outputs['fake_b']) assert outputs['fake_b'].size() == (1, 3, 256, 256) # test forward_train outputs = synthesizer(inputs, targets, [img_meta], test_mode=False) assert torch.equal(outputs['real_a'], data_batch['img_a']) assert torch.equal(outputs['real_b'], data_batch['img_b']) assert torch.is_tensor(outputs['fake_b']) assert outputs['fake_b'].size() == (1, 3, 256, 256) # test train_step outputs = synthesizer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['results'], dict) for v in [ 'loss_gan_d_fake', 'loss_gan_d_real', 'loss_gan_g', 'loss_pixel' ]: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['real_a'], data_batch['img_a']) assert torch.equal(outputs['results']['real_b'], data_batch['img_b']) assert torch.is_tensor(outputs['results']['fake_b']) assert outputs['results']['fake_b'].size() == (1, 3, 256, 256) # test train_step and forward_test (gpu) if torch.cuda.is_available(): synthesizer = synthesizer.cuda() optimizer = { 'generator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(synthesizer, 'generator').parameters())), 'discriminator': obj_from_dict( optim_cfg, torch.optim, dict( params=getattr(synthesizer, 'discriminator').parameters())) } data_batch_cuda = copy.deepcopy(data_batch) data_batch_cuda['img_a'] = inputs.cuda() data_batch_cuda['img_b'] = targets.cuda() data_batch_cuda['meta'] = [DC(img_meta, cpu_only=True).data] # forward_test with torch.no_grad(): outputs = synthesizer(data_batch_cuda['img_a'], data_batch_cuda['img_b'], data_batch_cuda['meta'], test_mode=True) assert torch.equal(outputs['real_a'], data_batch_cuda['img_a'].cpu()) assert torch.equal(outputs['real_b'], data_batch_cuda['img_b'].cpu()) assert torch.is_tensor(outputs['fake_b']) assert outputs['fake_b'].size() == (1, 3, 256, 256) # val_step with torch.no_grad(): outputs = synthesizer.val_step(data_batch_cuda) assert torch.equal(outputs['real_a'], data_batch_cuda['img_a'].cpu()) assert torch.equal(outputs['real_b'], data_batch_cuda['img_b'].cpu()) assert torch.is_tensor(outputs['fake_b']) assert outputs['fake_b'].size() == (1, 3, 256, 256) # test forward_train outputs = synthesizer(data_batch_cuda['img_a'], data_batch_cuda['img_b'], data_batch_cuda['meta'], test_mode=False) assert torch.equal(outputs['real_a'], data_batch_cuda['img_a']) assert torch.equal(outputs['real_b'], data_batch_cuda['img_b']) assert torch.is_tensor(outputs['fake_b']) assert outputs['fake_b'].size() == (1, 3, 256, 256) # train_step outputs = synthesizer.train_step(data_batch_cuda, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['results'], dict) for v in [ 'loss_gan_d_fake', 'loss_gan_d_real', 'loss_gan_g', 'loss_pixel' ]: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['real_a'], data_batch_cuda['img_a'].cpu()) assert torch.equal(outputs['results']['real_b'], data_batch_cuda['img_b'].cpu()) assert torch.is_tensor(outputs['results']['fake_b']) assert outputs['results']['fake_b'].size() == (1, 3, 256, 256) # test disc_steps and disc_init_steps data_batch['img_a'] = inputs.cpu() data_batch['img_b'] = targets.cpu() train_cfg = dict(disc_steps=2, disc_init_steps=2) synthesizer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) optimizer = { 'generator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(synthesizer, 'generator').parameters())), 'discriminator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(synthesizer, 'discriminator').parameters())) } # iter 0, 1 for i in range(2): assert synthesizer.step_counter == i outputs = synthesizer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['results'], dict) assert outputs['log_vars'].get('loss_gan_g') is None assert outputs['log_vars'].get('loss_pixel') is None for v in ['loss_gan_d_fake', 'loss_gan_d_real']: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['real_a'], data_batch['img_a']) assert torch.equal(outputs['results']['real_b'], data_batch['img_b']) assert torch.is_tensor(outputs['results']['fake_b']) assert outputs['results']['fake_b'].size() == (1, 3, 256, 256) assert synthesizer.step_counter == i + 1 # iter 2, 3, 4, 5 for i in range(2, 6): assert synthesizer.step_counter == i outputs = synthesizer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['results'], dict) log_check_list = [ 'loss_gan_d_fake', 'loss_gan_d_real', 'loss_gan_g', 'loss_pixel' ] if i % 2 == 1: assert outputs['log_vars'].get('loss_gan_g') is None assert outputs['log_vars'].get('loss_pixel') is None log_check_list.remove('loss_gan_g') log_check_list.remove('loss_pixel') for v in log_check_list: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['real_a'], data_batch['img_a']) assert torch.equal(outputs['results']['real_b'], data_batch['img_b']) assert torch.is_tensor(outputs['results']['fake_b']) assert outputs['results']['fake_b'].size() == (1, 3, 256, 256) assert synthesizer.step_counter == i + 1 # test without pixel loss model_cfg_ = copy.deepcopy(model_cfg) model_cfg_.pop('pixel_loss') synthesizer = build_model(model_cfg_, train_cfg=None, test_cfg=None) optimizer = { 'generator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(synthesizer, 'generator').parameters())), 'discriminator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(synthesizer, 'discriminator').parameters())) } data_batch['img_a'] = inputs.cpu() data_batch['img_b'] = targets.cpu() outputs = synthesizer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['results'], dict) assert outputs['log_vars'].get('loss_pixel') is None for v in ['loss_gan_d_fake', 'loss_gan_d_real', 'loss_gan_g']: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['real_a'], data_batch['img_a']) assert torch.equal(outputs['results']['real_b'], data_batch['img_b']) assert torch.is_tensor(outputs['results']['fake_b']) assert outputs['results']['fake_b'].size() == (1, 3, 256, 256) # test b2a translation data_batch['img_a'] = inputs.cpu() data_batch['img_b'] = targets.cpu() train_cfg = dict(direction='b2a') synthesizer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) optimizer = { 'generator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(synthesizer, 'generator').parameters())), 'discriminator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(synthesizer, 'discriminator').parameters())) } assert synthesizer.step_counter == 0 outputs = synthesizer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['results'], dict) for v in [ 'loss_gan_d_fake', 'loss_gan_d_real', 'loss_gan_g', 'loss_pixel' ]: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['real_a'], data_batch['img_b']) assert torch.equal(outputs['results']['real_b'], data_batch['img_a']) assert torch.is_tensor(outputs['results']['fake_b']) assert outputs['results']['fake_b'].size() == (1, 3, 256, 256) assert synthesizer.step_counter == 1 # test save image # show input train_cfg = None test_cfg = dict(show_input=True) synthesizer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) with patch.object(mmcv, 'imwrite', return_value=True): # test save path not None Assertion with pytest.raises(AssertionError): with torch.no_grad(): _ = synthesizer(inputs, targets, [img_meta], test_mode=True, save_image=True) # iteration is None with torch.no_grad(): outputs = synthesizer(inputs, targets, [img_meta], test_mode=True, save_image=True, save_path='save_path') assert torch.equal(outputs['real_a'], data_batch['img_a']) assert torch.equal(outputs['real_b'], data_batch['img_b']) assert torch.is_tensor(outputs['fake_b']) assert outputs['fake_b'].size() == (1, 3, 256, 256) assert outputs['saved_flag'] # iteration is not None with torch.no_grad(): outputs = synthesizer(inputs, targets, [img_meta], test_mode=True, save_image=True, save_path='save_path', iteration=1000) assert torch.equal(outputs['real_a'], data_batch['img_a']) assert torch.equal(outputs['real_b'], data_batch['img_b']) assert torch.is_tensor(outputs['fake_b']) assert outputs['fake_b'].size() == (1, 3, 256, 256) assert outputs['saved_flag'] # not show input train_cfg = None test_cfg = dict(show_input=False) synthesizer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) with patch.object(mmcv, 'imwrite', return_value=True): # test save path not None Assertion with pytest.raises(AssertionError): with torch.no_grad(): _ = synthesizer(inputs, targets, [img_meta], test_mode=True, save_image=True) # iteration is None with torch.no_grad(): outputs = synthesizer(inputs, targets, [img_meta], test_mode=True, save_image=True, save_path='save_path') assert torch.equal(outputs['real_a'], data_batch['img_a']) assert torch.equal(outputs['real_b'], data_batch['img_b']) assert torch.is_tensor(outputs['fake_b']) assert outputs['fake_b'].size() == (1, 3, 256, 256) assert outputs['saved_flag'] # iteration is not None with torch.no_grad(): outputs = synthesizer(inputs, targets, [img_meta], test_mode=True, save_image=True, save_path='save_path', iteration=1000) assert torch.equal(outputs['real_a'], data_batch['img_a']) assert torch.equal(outputs['real_b'], data_batch['img_b']) assert torch.is_tensor(outputs['fake_b']) assert outputs['fake_b'].size() == (1, 3, 256, 256) assert outputs['saved_flag']
def main(): args = parse_args() cfg = Config.fromfile(args.config) # set cudnn_benchmark if cfg.get('cudnn_benchmark', False): torch.backends.cudnn.benchmark = True # update configs according to CLI args if args.work_dir is not None: cfg.work_dir = args.work_dir if args.resume_from is not None: cfg.resume_from = args.resume_from cfg.gpus = args.gpus if args.autoscale_lr: # apply the linear scaling rule (https://arxiv.org/abs/1706.02677) cfg.optimizer['lr'] = cfg.optimizer['lr'] * cfg.gpus / 8 # init distributed env first, since logger depends on the dist info. if args.launcher == 'none': distributed = False else: distributed = True init_dist(args.launcher, **cfg.dist_params) # create work_dir mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir)) # init the logger before other steps timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime()) log_file = osp.join(cfg.work_dir, f'{timestamp}.log') logger = get_root_logger(log_file=log_file, log_level=cfg.log_level) # log env info env_info_dict = collect_env.collect_env() env_info = '\n'.join([f'{k}: {v}' for k, v in env_info_dict.items()]) dash_line = '-' * 60 + '\n' logger.info('Environment info:\n' + dash_line + env_info + '\n' + dash_line) # log some basic info logger.info('Distributed training: {}'.format(distributed)) logger.info('mmedit Version: {}'.format(__version__)) logger.info('Config:\n{}'.format(cfg.text)) # set random seeds if args.seed is not None: logger.info('Set random seed to {}, deterministic: {}'.format( args.seed, args.deterministic)) set_random_seed(args.seed, deterministic=args.deterministic) cfg.seed = args.seed model = build_model(cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg) datasets = [build_dataset(cfg.data.train)] if len(cfg.workflow) == 2: val_dataset = copy.deepcopy(cfg.data.val) val_dataset.pipeline = cfg.data.train.pipeline datasets.append(build_dataset(val_dataset)) if cfg.checkpoint_config is not None: # save version, config file content and class names in # checkpoints as meta data cfg.checkpoint_config.meta = dict( mmedit_version=__version__, config=cfg.text, ) # meta information meta = dict() if cfg.get('exp_name', None) is None: cfg['exp_name'] = osp.splitext(osp.basename(cfg.work_dir))[0] meta['exp_name'] = cfg.exp_name meta['mmedit Version'] = __version__ meta['seed'] = args.seed meta['env_info'] = env_info # add an attribute for visualization convenience train_model(model, datasets, cfg, distributed=distributed, validate=(not args.no_validate), timestamp=timestamp, meta=meta)
def test_basicvsr_model(): model_cfg = dict( type='BasicVSR', generator=dict(type='BasicVSRNet', mid_channels=64, num_blocks=30, spynet_pretrained=None), pixel_loss=dict(type='MSELoss', loss_weight=1.0, reduction='sum'), ) train_cfg = dict(fix_iter=1) train_cfg = mmcv.Config(train_cfg) test_cfg = None # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test attributes assert restorer.__class__.__name__ == 'BasicVSR' assert isinstance(restorer.generator, BasicVSRNet) assert isinstance(restorer.pixel_loss, MSELoss) # prepare data inputs = torch.rand(1, 5, 3, 64, 64) targets = torch.rand(1, 5, 3, 256, 256) if torch.cuda.is_available(): inputs = inputs.cuda() targets = targets.cuda() restorer = restorer.cuda() # prepare data and optimizer data_batch = {'lq': inputs, 'gt': targets} optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.9, 0.999)) optimizer = { 'generator': obj_from_dict(optim_cfg, torch.optim, dict(params=getattr(restorer, 'generator').parameters())) } # train_step (wihout updating spynet) outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['results']['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 5, 3, 256, 256) # train with spynet updated outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['results']['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 5, 3, 256, 256) # test forward_dummy with torch.no_grad(): output = restorer.forward_dummy(data_batch['lq']) assert torch.is_tensor(output) assert output.size() == (1, 5, 3, 256, 256) # forward_test with torch.no_grad(): outputs = restorer(**data_batch, test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 5, 3, 256, 256) with torch.no_grad(): outputs = restorer(inputs, test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 5, 3, 256, 256) # test with metric and save image train_cfg = mmcv.ConfigDict(fix_iter=1) test_cfg = dict(metrics=('PSNR', 'SSIM'), crop_border=0) test_cfg = mmcv.Config(test_cfg) data_batch = { 'lq': inputs, 'gt': targets, 'meta': [{ 'gt_path': 'fake_path/fake_name.png', 'key': '000' }] } restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) if torch.cuda.is_available(): restorer = restorer.cuda() with pytest.raises(AssertionError): # evaluation with metrics must have gt images restorer(lq=inputs, test_mode=True) with tempfile.TemporaryDirectory() as tmpdir: outputs = restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=None) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) outputs = restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=100) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) with pytest.raises(ValueError): # iteration should be number or None restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration='100')
def test_real_esrgan(): model_cfg = dict( type='RealESRGAN', generator=dict( type='MSRResNet', in_channels=3, out_channels=3, mid_channels=4, num_blocks=1, upscale_factor=4), discriminator=dict(type='ModifiedVGG', in_channels=3, mid_channels=2), pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean'), gan_loss=dict( type='GANLoss', gan_type='vanilla', loss_weight=1e-1, real_label_val=1.0, fake_label_val=0), is_use_sharpened_gt_in_pixel=True, is_use_sharpened_gt_in_percep=True, is_use_sharpened_gt_in_gan=True, is_use_ema=True, ) train_cfg = None test_cfg = None # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test attributes assert restorer.__class__.__name__ == 'RealESRGAN' assert isinstance(restorer.generator, MSRResNet) assert isinstance(restorer.discriminator, ModifiedVGG) assert isinstance(restorer.pixel_loss, L1Loss) assert isinstance(restorer.gan_loss, GANLoss) # prepare data inputs = torch.rand(1, 3, 32, 32) targets = torch.rand(1, 3, 128, 128) data_batch = {'lq': inputs, 'gt': targets, 'gt_unsharp': targets} # prepare optimizer optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.9, 0.999)) optimizer = { 'generator': obj_from_dict(optim_cfg, torch.optim, dict( params=getattr(restorer, 'generator').parameters())), 'discriminator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(restorer, 'discriminator').parameters())) } # no forward train in GAN models, raise ValueError with pytest.raises(ValueError): restorer(**data_batch, test_mode=False) # test forward_test data_batch.pop('gt_unsharp') with torch.no_grad(): outputs = restorer(**data_batch, test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq']) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 128, 128) # test forward_dummy with torch.no_grad(): output = restorer.forward_dummy(data_batch['lq']) assert torch.is_tensor(output) assert output.size() == (1, 3, 128, 128) # val_step with torch.no_grad(): outputs = restorer.val_step(data_batch) data_batch['gt_unsharp'] = targets assert torch.equal(outputs['lq'], data_batch['lq']) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 128, 128) # test train_step with patch.object( restorer, 'perceptual_loss', return_value=(torch.tensor(1.0), torch.tensor(2.0))): outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) for v in [ 'loss_perceptual', 'loss_gan', 'loss_d_real', 'loss_d_fake', 'loss_pix' ]: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq']) assert torch.equal(outputs['results']['gt'], data_batch['gt']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 128, 128) # test train_step and forward_test (gpu) if torch.cuda.is_available(): restorer = restorer.cuda() optimizer = { 'generator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(restorer, 'generator').parameters())), 'discriminator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(restorer, 'discriminator').parameters())) } data_batch = { 'lq': inputs.cuda(), 'gt': targets.cuda(), 'gt_unsharp': targets.cuda() } # forward_test data_batch.pop('gt_unsharp') with torch.no_grad(): outputs = restorer(**data_batch, test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 128, 128) # val_step with torch.no_grad(): outputs = restorer.val_step(data_batch) data_batch['gt_unsharp'] = targets.cuda() assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 128, 128) # train_step with patch.object( restorer, 'perceptual_loss', return_value=(torch.tensor(1.0).cuda(), torch.tensor(2.0).cuda())): outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) for v in [ 'loss_perceptual', 'loss_gan', 'loss_d_real', 'loss_d_fake', 'loss_pix' ]: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['results']['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 128, 128) # test disc_steps and disc_init_steps and start_iter data_batch = { 'lq': inputs.cpu(), 'gt': targets.cpu(), 'gt_unsharp': targets.cpu() } train_cfg = dict(disc_steps=2, disc_init_steps=2, start_iter=0) restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) with patch.object( restorer, 'perceptual_loss', return_value=(torch.tensor(1.0), torch.tensor(2.0))): outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) for v in ['loss_d_real', 'loss_d_fake']: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq']) assert torch.equal(outputs['results']['gt'], data_batch['gt']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 128, 128) # test no discriminator (testing mode) model_cfg_ = model_cfg.copy() model_cfg_.pop('discriminator') restorer = build_model(model_cfg_, train_cfg=train_cfg, test_cfg=test_cfg) data_batch.pop('gt_unsharp') with torch.no_grad(): outputs = restorer(**data_batch, test_mode=True) data_batch['gt_unsharp'] = targets.cpu() assert torch.equal(outputs['lq'], data_batch['lq']) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 128, 128) # test without pixel loss and perceptual loss model_cfg_ = model_cfg.copy() model_cfg_.pop('pixel_loss') restorer = build_model(model_cfg_, train_cfg=None, test_cfg=None) outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) for v in ['loss_gan', 'loss_d_real', 'loss_d_fake']: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq']) assert torch.equal(outputs['results']['gt'], data_batch['gt']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 128, 128) # test train_step w/o loss_percep restorer = build_model(model_cfg, train_cfg=None, test_cfg=None) with patch.object( restorer, 'perceptual_loss', return_value=(None, torch.tensor(2.0))): outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) for v in [ 'loss_style', 'loss_gan', 'loss_d_real', 'loss_d_fake', 'loss_pix' ]: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq']) assert torch.equal(outputs['results']['gt'], data_batch['gt']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 128, 128) # test train_step w/o loss_style restorer = build_model(model_cfg, train_cfg=None, test_cfg=None) with patch.object( restorer, 'perceptual_loss', return_value=(torch.tensor(2.0), None)): outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) for v in [ 'loss_perceptual', 'loss_gan', 'loss_d_real', 'loss_d_fake', 'loss_pix' ]: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq']) assert torch.equal(outputs['results']['gt'], data_batch['gt']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 128, 128)
if model_type == 'mattor' and args.trimap_path is None: raise ValueError('Please set `--trimap-path` to convert mattor model.') assert args.opset_version == 11, 'MMEditing only support opset 11 now' config = mmcv.Config.fromfile(args.config) config.model.pretrained = None # ONNX does not support spectral norm if model_type == 'mattor': if hasattr(config.model.backbone.encoder, 'with_spectral_norm'): config.model.backbone.encoder.with_spectral_norm = False config.model.backbone.decoder.with_spectral_norm = False config.test_cfg.metrics = None # build the model model = build_model(config.model, test_cfg=config.test_cfg) checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu') # remove alpha from test_pipeline if model_type == 'mattor': keys_to_remove = ['alpha', 'ori_alpha'] elif model_type == 'restorer': keys_to_remove = ['gt', 'gt_path'] for key in keys_to_remove: for pipeline in list(config.test_pipeline): if 'key' in pipeline and key == pipeline['key']: config.test_pipeline.remove(pipeline) if 'keys' in pipeline and key in pipeline['keys']: pipeline['keys'].remove(key) if len(pipeline['keys']) == 0: config.test_pipeline.remove(pipeline)
def test_tdan_model(): model_cfg = dict( type='TDAN', generator=dict(type='TDANNet', in_channels=3, mid_channels=64, out_channels=3, num_blocks_before_align=5, num_blocks_after_align=10), pixel_loss=dict(type='MSELoss', loss_weight=1.0, reduction='sum'), lq_pixel_loss=dict(type='MSELoss', loss_weight=1.0, reduction='sum'), ) train_cfg = None test_cfg = None # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test attributes assert restorer.__class__.__name__ == 'TDAN' assert isinstance(restorer.generator, TDANNet) assert isinstance(restorer.pixel_loss, MSELoss) # prepare data inputs = torch.rand(1, 5, 3, 8, 8) targets = torch.rand(1, 3, 32, 32) # test train_step and forward_test (gpu) if torch.cuda.is_available(): restorer = restorer.cuda() data_batch = {'lq': inputs.cuda(), 'gt': targets.cuda()} # prepare optimizer optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.9, 0.999)) optimizer = { 'generator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(restorer, 'generator').parameters())) } # train_step outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['results']['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 32, 32) # test forward_dummy with torch.no_grad(): output = restorer.forward_dummy(data_batch['lq']) assert isinstance(output, tuple) assert torch.is_tensor(output[0]) assert output[0].size() == (1, 3, 32, 32) assert torch.is_tensor(output[1]) assert output[1].size() == (1, 5, 3, 8, 8) # forward_test with torch.no_grad(): outputs = restorer(**data_batch, test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 32, 32) with torch.no_grad(): outputs = restorer(inputs.cuda(), test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 32, 32) # test with metric and save image if torch.cuda.is_available(): train_cfg = mmcv.ConfigDict(tsa_iter=1) test_cfg = dict(metrics=('PSNR', 'SSIM'), crop_border=0) test_cfg = mmcv.Config(test_cfg) data_batch = { 'lq': inputs.cuda(), 'gt': targets.cuda(), 'meta': [{ 'gt_path': 'fake_path/fake_name.png', 'key': '000/00000000' }] } restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg).cuda() with pytest.raises(AssertionError): # evaluation with metrics must have gt images restorer(lq=inputs.cuda(), test_mode=True) with tempfile.TemporaryDirectory() as tmpdir: outputs = restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=None) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) outputs = restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=100) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) with pytest.raises(ValueError): # iteration should be number or None restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration='100')
def test_basic_restorer(): model_cfg = dict(type='BasicRestorer', generator=dict(type='MSRResNet', in_channels=3, out_channels=3, mid_channels=4, num_blocks=1, upscale_factor=4), pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean')) train_cfg = None test_cfg = None # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test attributes assert restorer.__class__.__name__ == 'BasicRestorer' assert isinstance(restorer.generator, MSRResNet) assert isinstance(restorer.pixel_loss, L1Loss) # prepare data inputs = torch.rand(1, 3, 2, 2) targets = torch.rand(1, 3, 8, 8) data_batch = {'lq': inputs, 'gt': targets} # prepare optimizer optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.9, 0.999)) optimizer = { 'generator': obj_from_dict(optim_cfg, torch.optim, dict(params=restorer.parameters())) } # test forward train outputs = restorer(**data_batch, test_mode=False) assert isinstance(outputs, dict) assert isinstance(outputs['losses'], dict) assert isinstance(outputs['losses']['loss_pix'], torch.FloatTensor) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq']) assert torch.equal(outputs['results']['gt'], data_batch['gt']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 8, 8) # test forward_test with torch.no_grad(): outputs = restorer(**data_batch, test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq']) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 8, 8) # test forward_dummy with torch.no_grad(): output = restorer.forward_dummy(data_batch['lq']) assert torch.is_tensor(output) assert output.size() == (1, 3, 8, 8) # test train_step outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq']) assert torch.equal(outputs['results']['gt'], data_batch['gt']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 8, 8) # test train_step and forward_test (gpu) if torch.cuda.is_available(): restorer = restorer.cuda() optimizer['generator'] = obj_from_dict( optim_cfg, torch.optim, dict(params=restorer.parameters())) data_batch = {'lq': inputs.cuda(), 'gt': targets.cuda()} # test forward train outputs = restorer(**data_batch, test_mode=False) assert isinstance(outputs, dict) assert isinstance(outputs['losses'], dict) assert isinstance(outputs['losses']['loss_pix'], torch.cuda.FloatTensor) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['results']['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 8, 8) # forward_test with torch.no_grad(): outputs = restorer(**data_batch, test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 8, 8) # train_step outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['results']['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 8, 8) # test with metric and save image test_cfg = dict(metrics=('PSNR', 'SSIM'), crop_border=0) test_cfg = mmcv.Config(test_cfg) data_batch = { 'lq': inputs, 'gt': targets, 'meta': [{ 'lq_path': 'fake_path/fake_name.png' }] } restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) with pytest.raises(AssertionError): # evaluation with metrics must have gt images restorer(lq=inputs, test_mode=True) with tempfile.TemporaryDirectory() as tmpdir: outputs = restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=None) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) outputs = restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=100) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) with pytest.raises(ValueError): # iteration should be number or None restorer(**data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration='100')
def test_real_basicvsr(): model_cfg = dict( type='RealBasicVSR', generator=dict(type='RealBasicVSRNet'), discriminator=dict(type='UNetDiscriminatorWithSpectralNorm', in_channels=3, mid_channels=64, skip_connection=True), pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean'), cleaning_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean'), gan_loss=dict(type='GANLoss', gan_type='vanilla', loss_weight=1e-1, real_label_val=1.0, fake_label_val=0), is_use_sharpened_gt_in_pixel=True, is_use_sharpened_gt_in_percep=True, is_use_sharpened_gt_in_gan=True, is_use_ema=True, ) train_cfg = None test_cfg = None # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test attributes assert restorer.__class__.__name__ == 'RealBasicVSR' assert isinstance(restorer.generator, RealBasicVSRNet) assert isinstance(restorer.discriminator, UNetDiscriminatorWithSpectralNorm) assert isinstance(restorer.pixel_loss, L1Loss) assert isinstance(restorer.gan_loss, GANLoss) # prepare data inputs = torch.rand(1, 5, 3, 64, 64) targets = torch.rand(1, 5, 3, 256, 256) data_batch = {'lq': inputs, 'gt': targets, 'gt_unsharp': targets} # prepare optimizer optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.9, 0.999)) optimizer = { 'generator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(restorer, 'generator').parameters())), 'discriminator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(restorer, 'discriminator').parameters())) } # no forward train in GAN models, raise ValueError with pytest.raises(ValueError): restorer(**data_batch, test_mode=False) # test train_step with patch.object(restorer, 'perceptual_loss', return_value=(torch.tensor(1.0), torch.tensor(2.0))): outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) for v in [ 'loss_perceptual', 'loss_gan', 'loss_d_real', 'loss_d_fake', 'loss_pix', 'loss_clean' ]: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq']) assert torch.equal(outputs['results']['gt'], data_batch['gt']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (5, 3, 256, 256) # test train_step (gpu) if torch.cuda.is_available(): restorer = restorer.cuda() optimizer = { 'generator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(restorer, 'generator').parameters())), 'discriminator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(restorer, 'discriminator').parameters())) } data_batch = { 'lq': inputs.cuda(), 'gt': targets.cuda(), 'gt_unsharp': targets.cuda() } # train_step with patch.object(restorer, 'perceptual_loss', return_value=(torch.tensor(1.0).cuda(), torch.tensor(2.0).cuda())): outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) for v in [ 'loss_perceptual', 'loss_gan', 'loss_d_real', 'loss_d_fake', 'loss_pix', 'loss_clean' ]: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['results']['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (5, 3, 256, 256) # test disc_steps and disc_init_steps and start_iter data_batch = { 'lq': inputs.cpu(), 'gt': targets.cpu(), 'gt_unsharp': targets.cpu() } train_cfg = dict(disc_steps=2, disc_init_steps=2, start_iter=0) restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) with patch.object(restorer, 'perceptual_loss', return_value=(torch.tensor(1.0), torch.tensor(2.0))): outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) for v in ['loss_d_real', 'loss_d_fake']: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq']) assert torch.equal(outputs['results']['gt'], data_batch['gt']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (5, 3, 256, 256) # test without pixel loss and perceptual loss model_cfg_ = model_cfg.copy() model_cfg_.pop('pixel_loss') restorer = build_model(model_cfg_, train_cfg=None, test_cfg=None) outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) for v in ['loss_gan', 'loss_d_real', 'loss_d_fake']: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq']) assert torch.equal(outputs['results']['gt'], data_batch['gt']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (5, 3, 256, 256) # test train_step w/o loss_percep restorer = build_model(model_cfg, train_cfg=None, test_cfg=None) with patch.object(restorer, 'perceptual_loss', return_value=(None, torch.tensor(2.0))): outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) for v in [ 'loss_style', 'loss_gan', 'loss_d_real', 'loss_d_fake', 'loss_pix', 'loss_clean' ]: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq']) assert torch.equal(outputs['results']['gt'], data_batch['gt']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (5, 3, 256, 256) # test train_step w/o loss_style restorer = build_model(model_cfg, train_cfg=None, test_cfg=None) with patch.object(restorer, 'perceptual_loss', return_value=(torch.tensor(2.0), None)): outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) for v in [ 'loss_perceptual', 'loss_gan', 'loss_d_real', 'loss_d_fake', 'loss_pix', 'loss_clean' ]: assert isinstance(outputs['log_vars'][v], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq']) assert torch.equal(outputs['results']['gt'], data_batch['gt']) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (5, 3, 256, 256)
def test_dim(): model_cfg, train_cfg, test_cfg = _get_model_cfg( 'dim/dim_stage3_v16_pln_1x1_1000k_comp1k.py') model_cfg['pretrained'] = None # 1. test dim model with refiner train_cfg.train_refiner = True test_cfg.refine = True # test model forward in train mode model = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) input_train = _demo_input_train((64, 64)) output_train = model(**input_train) assert output_train['num_samples'] == 1 assert_dict_keys_equal(output_train['losses'], ['loss_alpha', 'loss_comp', 'loss_refine']) # test model forward in train mode with gpu if torch.cuda.is_available(): model = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) model.cuda() input_train = _demo_input_train((64, 64), cuda=True) output_train = model(**input_train) assert output_train['num_samples'] == 1 assert_dict_keys_equal(output_train['losses'], ['loss_alpha', 'loss_comp', 'loss_refine']) # test model forward in test mode with torch.no_grad(): model = build_model(model_cfg, train_cfg=None, test_cfg=test_cfg) input_test = _demo_input_test((64, 64)) output_test = model(**input_test, test_mode=True) assert isinstance(output_test['pred_alpha'], np.ndarray) assert_dict_keys_equal(output_test['eval_result'], ['SAD', 'MSE', 'GRAD', 'CONN']) # test model forward in test mode with gpu if torch.cuda.is_available(): model = build_model(model_cfg, train_cfg=None, test_cfg=test_cfg) model.cuda() input_test = _demo_input_test((64, 64), cuda=True) output_test = model(**input_test, test_mode=True) assert isinstance(output_test['pred_alpha'], np.ndarray) assert_dict_keys_equal(output_test['eval_result'], ['SAD', 'MSE', 'GRAD', 'CONN']) # 2. test dim model without refiner model_cfg['refiner'] = None test_cfg['metrics'] = None # test model forward in train mode model = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) input_train = _demo_input_train((64, 64)) output_train = model(**input_train) assert output_train['num_samples'] == 1 assert_dict_keys_equal(output_train['losses'], ['loss_alpha', 'loss_comp']) # test model forward in train mode with gpu if torch.cuda.is_available(): model = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) model.cuda() input_train = _demo_input_train((64, 64), cuda=True) output_train = model(**input_train) assert output_train['num_samples'] == 1 assert_dict_keys_equal(output_train['losses'], ['loss_alpha', 'loss_comp']) # test model forward in test mode with torch.no_grad(): model = build_model(model_cfg, train_cfg=None, test_cfg=test_cfg) input_test = _demo_input_test((64, 64)) output_test = model(**input_test, test_mode=True) assert isinstance(output_test['pred_alpha'], np.ndarray) assert output_test['eval_result'] is None # check test with gpu if torch.cuda.is_available(): model = build_model(model_cfg, train_cfg=None, test_cfg=test_cfg) model.cuda() input_test = _demo_input_test((64, 64), cuda=True) output_test = model(**input_test, test_mode=True) assert isinstance(output_test['pred_alpha'], np.ndarray) assert output_test['eval_result'] is None # test forward_dummy model.cpu().eval() inputs = torch.ones((1, 4, 32, 32)) model.forward_dummy(inputs)
def test_ttsr(): model_cfg = dict(type='TTSR', generator=dict(type='TTSRNet', in_channels=3, out_channels=3, mid_channels=64, num_blocks=(16, 16, 8, 4)), extractor=dict(type='LTE'), transformer=dict(type='SearchTransformer'), discriminator=dict(type='TTSRDiscriminator', in_size=64), pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean'), perceptual_loss=dict(type='PerceptualLoss', layer_weights={'29': 1.0}, vgg_type='vgg19', perceptual_weight=1e-2, style_weight=0.001, criterion='mse'), transferal_perceptual_loss=dict( type='TransferalPerceptualLoss', loss_weight=1e-2, use_attention=False, criterion='mse'), gan_loss=dict(type='GANLoss', gan_type='vanilla', loss_weight=1e-3, real_label_val=1.0, fake_label_val=0)) scale = 4 train_cfg = None test_cfg = Config(dict(metrics=['PSNR', 'SSIM'], crop_border=scale)) # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) model_cfg = dict(type='TTSR', generator=dict(type='TTSRNet', in_channels=3, out_channels=3, mid_channels=64, num_blocks=(16, 16, 8, 4)), extractor=dict(type='LTE'), transformer=dict(type='SearchTransformer'), discriminator=dict(type='TTSRDiscriminator', in_size=64), pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='mean'), perceptual_loss=dict(type='PerceptualLoss', layer_weights={'29': 1.0}, vgg_type='vgg19', perceptual_weight=1e-2, style_weight=0.001, criterion='mse'), transferal_perceptual_loss=dict( type='TransferalPerceptualLoss', loss_weight=1e-2, use_attention=False, criterion='mse'), gan_loss=dict(type='GANLoss', gan_type='vanilla', loss_weight=1e-3, real_label_val=1.0, fake_label_val=0)) scale = 4 train_cfg = None test_cfg = Config(dict(metrics=['PSNR', 'SSIM'], crop_border=scale)) # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test attributes assert restorer.__class__.__name__ == 'TTSR' # prepare data inputs = torch.rand(1, 3, 16, 16) targets = torch.rand(1, 3, 64, 64) ref = torch.rand(1, 3, 64, 64) data_batch = { 'lq': inputs, 'gt': targets, 'ref': ref, 'lq_up': ref, 'ref_downup': ref } # prepare optimizer optim_cfg_g = dict(type='Adam', lr=1e-4, betas=(0.9, 0.999)) optim_cfg_d = dict(type='Adam', lr=1e-4, betas=(0.9, 0.999)) optimizer = dict( generator=obj_from_dict(optim_cfg_g, torch.optim, dict(params=restorer.parameters())), discriminator=obj_from_dict(optim_cfg_d, torch.optim, dict(params=restorer.parameters()))) # test train_step and forward_test (cpu) outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert outputs['results']['lq'].shape == data_batch['lq'].shape assert outputs['results']['gt'].shape == data_batch['gt'].shape assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 64, 64) # test train_step and forward_test (gpu) if torch.cuda.is_available(): restorer = restorer.cuda() data_batch = { 'lq': inputs.cuda(), 'gt': targets.cuda(), 'ref': ref.cuda(), 'lq_up': ref.cuda(), 'ref_downup': ref.cuda() } # train_step optimizer = dict( generator=obj_from_dict(optim_cfg_g, torch.optim, dict(params=restorer.parameters())), discriminator=obj_from_dict(optim_cfg_d, torch.optim, dict(params=restorer.parameters()))) outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert outputs['results']['lq'].shape == data_batch['lq'].shape assert outputs['results']['gt'].shape == data_batch['gt'].shape assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 64, 64) # val_step result = restorer.val_step(data_batch, meta=[{'gt_path': ''}]) assert isinstance(result, dict) assert isinstance(result['eval_result'], dict) assert result['eval_result'].keys() == set({'PSNR', 'SSIM'}) assert isinstance(result['eval_result']['PSNR'], np.float64) assert isinstance(result['eval_result']['SSIM'], np.float64)
def test_one_stage_inpaintor(): cfg = Config.fromfile('tests/data/inpaintor_config/one_stage_gl.py') # mock perceptual loss for test speed cfg.model.loss_composed_percep = None inpaintor = build_model(cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg) # modify attributes for mocking inpaintor.with_composed_percep_loss = True inpaintor.loss_percep = None # test attributes assert inpaintor.__class__.__name__ == 'OneStageInpaintor' assert isinstance(inpaintor.generator, GLEncoderDecoder) assert inpaintor.with_l1_hole_loss assert inpaintor.with_l1_valid_loss assert inpaintor.with_tv_loss assert inpaintor.with_composed_percep_loss assert inpaintor.with_out_percep_loss assert inpaintor.with_gan assert inpaintor.with_gp_loss assert inpaintor.with_disc_shift_loss assert inpaintor.is_train assert inpaintor.train_cfg['disc_step'] == 1 assert inpaintor.disc_step_count == 0 with patch.object(inpaintor, 'loss_percep', return_value=(torch.tensor(1.0), None)): input_x = torch.randn(1, 3, 256, 256) with pytest.raises(NotImplementedError): inpaintor.forward_train(input_x) if torch.cuda.is_available(): gt_img = torch.randn(1, 3, 256, 256).cuda() mask = torch.zeros_like(gt_img)[:, 0:1, ...] mask[..., 20:100, 100:120] = 1. masked_img = gt_img * (1. - mask) inpaintor.cuda() data_batch = dict(gt_img=gt_img, mask=mask, masked_img=masked_img) output = inpaintor.forward_test(**data_batch) assert 'eval_result' in output output = inpaintor.val_step(data_batch) assert 'eval_result' in output optim_g = torch.optim.SGD(inpaintor.generator.parameters(), lr=0.1) optim_d = torch.optim.SGD(inpaintor.disc.parameters(), lr=0.1) optim_dict = dict(generator=optim_g, disc=optim_d) outputs = inpaintor.train_step(data_batch, optim_dict) assert outputs['num_samples'] == 1 results = outputs['results'] assert results['fake_res'].shape == (1, 3, 256, 256) assert 'loss_l1_hole' in outputs['log_vars'] assert 'loss_l1_valid' in outputs['log_vars'] assert 'loss_composed_percep' in outputs['log_vars'] assert 'loss_composed_style' not in outputs['log_vars'] assert 'loss_out_percep' in outputs['log_vars'] assert 'loss_out_style' not in outputs['log_vars'] assert 'loss_tv' in outputs['log_vars'] assert 'fake_loss' in outputs['log_vars'] assert 'real_loss' in outputs['log_vars'] assert 'loss_g_fake' in outputs['log_vars'] # test forward dummy res = inpaintor.forward_dummy(torch.cat([masked_img, mask], dim=1)) assert res.shape == (1, 3, 256, 256) # test forward test w/o save image outputs = inpaintor.forward_test(masked_img[0:1], mask[0:1], gt_img=gt_img[0:1, ...]) assert 'eval_result' in outputs assert outputs['eval_result']['l1'] > 0 assert outputs['eval_result']['psnr'] > 0 assert outputs['eval_result']['ssim'] > 0 # test forward test w/o eval metrics inpaintor.test_cfg = dict() inpaintor.eval_with_metrics = False outputs = inpaintor.forward_test(masked_img[0:1], mask[0:1]) for key in ['fake_res', 'fake_img']: assert outputs[key].size() == (1, 3, 256, 256) # test forward test w/ save image with tempfile.TemporaryDirectory() as tmpdir: outputs = inpaintor.forward_test( masked_img[0:1], mask[0:1], save_image=True, save_path=tmpdir, iteration=4396, meta=[dict(gt_img_path='igccc.png')]) assert os.path.exists(os.path.join(tmpdir, 'igccc_4396.png')) # test forward test w/ save image w/ gt_img with tempfile.TemporaryDirectory() as tmpdir: outputs = inpaintor.forward_test( masked_img[0:1], mask[0:1], save_image=True, save_path=tmpdir, meta=[dict(gt_img_path='igccc.png')], gt_img=gt_img[0:1, ...]) assert os.path.exists(os.path.join(tmpdir, 'igccc.png')) with pytest.raises(AssertionError): outputs = inpaintor.forward_test(masked_img[0:1], mask[0:1], save_image=True, save_path=tmpdir, iteration=4396, gt_img=gt_img[0:1, ...]) with pytest.raises(AssertionError): outputs = inpaintor.forward_test( masked_img[0:1], mask[0:1], save_image=True, save_path=None, iteration=4396, meta=[dict(gt_img_path='igccc.png')], gt_img=gt_img[0:1, ...]) cfg_ = copy.deepcopy(cfg) cfg_.train_cfg.disc_step = 2 inpaintor = build_model(cfg_.model, train_cfg=cfg_.train_cfg, test_cfg=cfg_.test_cfg) inpaintor.cuda() assert inpaintor.train_cfg.disc_step == 2 outputs = inpaintor.train_step(data_batch, optim_dict) assert 'loss_l1_hole' not in outputs['log_vars']
def main(): args = parse_args() cfg = mmcv.Config.fromfile(args.config) # set cudnn_benchmark if cfg.get('cudnn_benchmark', False): torch.backends.cudnn.benchmark = True cfg.model.pretrained = None # init distributed env first, since logger depends on the dist info. if args.launcher == 'none': distributed = False else: distributed = True init_dist(args.launcher, **cfg.dist_params) rank, _ = get_dist_info() # set random seeds if args.seed is not None: if rank == 0: print('set random seed to', args.seed) set_random_seed(args.seed, deterministic=args.deterministic) # build the dataloader # TODO: support multiple images per gpu (only minor changes are needed) dataset = build_dataset(cfg.data.test) loader_cfg = { **dict((k, cfg.data[k]) for k in ['workers_per_gpu'] if k in cfg.data), **dict(samples_per_gpu=1, drop_last=False, shuffle=False, dist=distributed), **cfg.data.get('test_dataloader', {}) } data_loader = build_dataloader(dataset, **loader_cfg) # build the model and load checkpoint model = build_model(cfg.model, train_cfg=None, test_cfg=cfg.test_cfg) args.save_image = args.save_path is not None empty_cache = cfg.get('empty_cache', False) if not distributed: _ = load_checkpoint(model, args.checkpoint, map_location='cpu') model = MMDataParallel(model, device_ids=[0]) outputs = single_gpu_test(model, data_loader, save_path=args.save_path, save_image=args.save_image) else: find_unused_parameters = cfg.get('find_unused_parameters', False) model = DistributedDataParallelWrapper( model, device_ids=[torch.cuda.current_device()], broadcast_buffers=False, find_unused_parameters=find_unused_parameters) device_id = torch.cuda.current_device() _ = load_checkpoint( model, args.checkpoint, map_location=lambda storage, loc: storage.cuda(device_id)) outputs = multi_gpu_test(model, data_loader, args.tmpdir, args.gpu_collect, save_path=args.save_path, save_image=args.save_image, empty_cache=empty_cache) if rank == 0: print('') # print metrics stats = dataset.evaluate(outputs) for stat in stats: print('Eval-{}: {}'.format(stat, stats[stat])) # save result pickle if args.out: print('writing results to {}'.format(args.out)) mmcv.dump(outputs, args.out)
def test_edvr_model(): model_cfg = dict( type='EDVR', generator=dict( type='EDVRNet', in_channels=3, out_channels=3, mid_channels=8, num_frames=5, deform_groups=2, num_blocks_extraction=1, num_blocks_reconstruction=1, center_frame_idx=2, with_tsa=False), pixel_loss=dict(type='L1Loss', loss_weight=1.0, reduction='sum'), ) train_cfg = None test_cfg = None # build restorer restorer = build_model(model_cfg, train_cfg=train_cfg, test_cfg=test_cfg) # test attributes assert restorer.__class__.__name__ == 'EDVR' assert isinstance(restorer.generator, EDVRNet) assert isinstance(restorer.pixel_loss, L1Loss) # prepare data inputs = torch.rand(1, 5, 3, 8, 8) targets = torch.rand(1, 3, 32, 32) # test train_step and forward_test (gpu) if torch.cuda.is_available(): restorer = restorer.cuda() data_batch = {'lq': inputs.cuda(), 'gt': targets.cuda()} # prepare optimizer optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.9, 0.999)) optimizer = { 'generator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(restorer, 'generator').parameters())) } # train_step outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['results']['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 32, 32) # with TSA model_cfg['generator']['with_tsa'] = True with pytest.raises(KeyError): # In TSA mode, train_cfg must contain "tsa_iter" train_cfg = dict(other_conent='xxx') restorer = build_model( model_cfg, train_cfg=train_cfg, test_cfg=test_cfg).cuda() outputs = restorer.train_step(data_batch, optimizer) train_cfg = None restorer = build_model( model_cfg, train_cfg=train_cfg, test_cfg=test_cfg).cuda() outputs = restorer.train_step(data_batch, optimizer) train_cfg = mmcv.ConfigDict(tsa_iter=1) restorer = build_model( model_cfg, train_cfg=train_cfg, test_cfg=test_cfg).cuda() optimizer = { 'generator': obj_from_dict( optim_cfg, torch.optim, dict(params=getattr(restorer, 'generator').parameters())) } # train without updating tsa module outputs = restorer.train_step(data_batch, optimizer) # train with updating tsa module outputs = restorer.train_step(data_batch, optimizer) assert isinstance(outputs, dict) assert isinstance(outputs['log_vars'], dict) assert isinstance(outputs['log_vars']['loss_pix'], float) assert outputs['num_samples'] == 1 assert torch.equal(outputs['results']['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['results']['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['results']['output']) assert outputs['results']['output'].size() == (1, 3, 32, 32) # test forward_dummy with torch.no_grad(): output = restorer.forward_dummy(data_batch['lq']) assert torch.is_tensor(output) assert output.size() == (1, 3, 32, 32) # forward_test with torch.no_grad(): outputs = restorer(**data_batch, test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.equal(outputs['gt'], data_batch['gt'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 32, 32) with torch.no_grad(): outputs = restorer(inputs.cuda(), test_mode=True) assert torch.equal(outputs['lq'], data_batch['lq'].cpu()) assert torch.is_tensor(outputs['output']) assert outputs['output'].size() == (1, 3, 32, 32) # test with metric and save image if torch.cuda.is_available(): train_cfg = mmcv.ConfigDict(tsa_iter=1) test_cfg = dict(metrics=('PSNR', 'SSIM'), crop_border=0) test_cfg = mmcv.Config(test_cfg) data_batch = { 'lq': inputs.cuda(), 'gt': targets.cuda(), 'meta': [{ 'gt_path': 'fake_path/fake_name.png', 'key': '000/00000000' }] } restorer = build_model( model_cfg, train_cfg=train_cfg, test_cfg=test_cfg).cuda() with pytest.raises(AssertionError): # evaluation with metrics must have gt images restorer(lq=inputs.cuda(), test_mode=True) with tempfile.TemporaryDirectory() as tmpdir: outputs = restorer( **data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=None) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) outputs = restorer( **data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration=100) assert isinstance(outputs, dict) assert isinstance(outputs['eval_result'], dict) assert isinstance(outputs['eval_result']['PSNR'], float) assert isinstance(outputs['eval_result']['SSIM'], float) with pytest.raises(ValueError): # iteration should be number or None restorer( **data_batch, test_mode=True, save_image=True, save_path=tmpdir, iteration='100')