def test_gl_encdec(): input_x = torch.randn(1, 4, 256, 256) template_cfg = dict(type='AOTEncoderDecoder') aot_encdec = build_backbone(template_cfg) aot_encdec.init_weights() output = aot_encdec(input_x) assert output.shape == (1, 3, 256, 256) cfg_ = template_cfg.copy() cfg_['encoder'] = dict(type='AOTEncoder') aot_encdec = build_backbone(cfg_) output = aot_encdec(input_x) assert output.shape == (1, 3, 256, 256) cfg_ = template_cfg.copy() cfg_['decoder'] = dict(type='AOTDecoder') aot_encdec = build_backbone(cfg_) output = aot_encdec(input_x) assert output.shape == (1, 3, 256, 256) if torch.cuda.is_available(): aot_encdec = build_backbone(template_cfg) aot_encdec.init_weights() aot_encdec = aot_encdec.cuda() output = aot_encdec(input_x.cuda()) assert output.shape == (1, 3, 256, 256)
def test_gl_encdec(): input_x = torch.randn(1, 4, 256, 256) template_cfg = dict(type='GLEncoderDecoder') gl_encdec = build_backbone(template_cfg) gl_encdec.init_weights() output = gl_encdec(input_x) assert output.shape == (1, 3, 256, 256) cfg_ = template_cfg.copy() cfg_['decoder'] = dict(type='GLDecoder', out_act='sigmoid') gl_encdec = build_backbone(cfg_) output = gl_encdec(input_x) assert output.shape == (1, 3, 256, 256) with pytest.raises(ValueError): cfg_ = template_cfg.copy() cfg_['decoder'] = dict(type='GLDecoder', out_act='igccc') gl_encdec = build_backbone(cfg_) with pytest.raises(TypeError): gl_encdec.init_weights(pretrained=dict(igccc=4396)) if torch.cuda.is_available(): gl_encdec = build_backbone(template_cfg) gl_encdec.init_weights() gl_encdec = gl_encdec.cuda() output = gl_encdec(input_x.cuda()) assert output.shape == (1, 3, 256, 256)
def test_dic_net(): model_cfg = dict( type='DICNet', in_channels=3, out_channels=3, mid_channels=48, num_blocks=6, hg_mid_channels=256, hg_num_keypoints=68, num_steps=4, upscale_factor=8, detach_attention=False) # build model model = build_backbone(model_cfg) # test attributes assert model.__class__.__name__ == 'DICNet' # prepare data inputs = torch.rand(1, 3, 16, 16) targets = torch.rand(1, 3, 128, 128) # prepare loss loss_function = nn.L1Loss() # prepare optimizer optimizer = torch.optim.Adam(model.parameters()) # test on cpu output, _ = model(inputs) optimizer.zero_grad() loss = loss_function(output[-1], targets) loss.backward() optimizer.step() assert len(output) == 4 assert torch.is_tensor(output[-1]) assert output[-1].shape == targets.shape # test on gpu if torch.cuda.is_available(): model = model.cuda() optimizer = torch.optim.Adam(model.parameters()) inputs = inputs.cuda() targets = targets.cuda() output, _ = model(inputs) optimizer.zero_grad() loss = loss_function(output[-1], targets) loss.backward() optimizer.step() assert len(output) == 4 assert torch.is_tensor(output[-1]) assert output[-1].shape == targets.shape with pytest.raises(OSError): model.init_weights('') with pytest.raises(TypeError): model.init_weights(1)
def test_rdn(): scale = 4 model_cfg = dict( type='RDN', in_channels=3, out_channels=3, mid_channels=64, num_blocks=16, upscale_factor=scale) # build model model = build_backbone(model_cfg) # test attributes assert model.__class__.__name__ == 'RDN' # prepare data inputs = torch.rand(1, 3, 32, 16) targets = torch.rand(1, 3, 128, 64) # prepare loss loss_function = nn.L1Loss() # prepare optimizer optimizer = torch.optim.Adam(model.parameters()) # test on cpu output = model(inputs) optimizer.zero_grad() loss = loss_function(output, targets) loss.backward() optimizer.step() assert torch.is_tensor(output) assert output.shape == targets.shape # test on gpu if torch.cuda.is_available(): model = model.cuda() optimizer = torch.optim.Adam(model.parameters()) inputs = inputs.cuda() targets = targets.cuda() output = model(inputs) optimizer.zero_grad() loss = loss_function(output, targets) loss.backward() optimizer.step() assert torch.is_tensor(output) assert output.shape == targets.shape
def test_liif_rdn(): model_cfg = dict(type='LIIFRDN', encoder=dict(type='RDN', in_channels=3, out_channels=3, mid_channels=64, num_blocks=16, upscale_factor=4, num_layers=8, channel_growth=64), 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) # build model model = build_backbone(model_cfg) # test attributes assert model.__class__.__name__ == 'LIIFRDN' # 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) # test on cpu output = model(inputs, coord, cell) output = model(inputs, coord, cell, True) assert torch.is_tensor(output) assert output.shape == targets.shape # test on gpu if torch.cuda.is_available(): model = model.cuda() inputs = inputs.cuda() targets = targets.cuda() coord = coord.cuda() cell = cell.cuda() output = model(inputs, coord, cell) output = model(inputs, coord, cell, True) assert torch.is_tensor(output) assert output.shape == targets.shape
def test_ttsr_net(): inputs = torch.rand(2, 3, 24, 24) soft_attention = torch.rand(2, 1, 24, 24) t_level3 = torch.rand(2, 64, 24, 24) t_level2 = torch.rand(2, 32, 48, 48) t_level1 = torch.rand(2, 16, 96, 96) ttsr_cfg = dict(type='TTSRNet', in_channels=3, out_channels=3, mid_channels=16, texture_channels=16) ttsr = build_backbone(ttsr_cfg) outputs = ttsr(inputs, soft_attention, (t_level3, t_level2, t_level1)) assert outputs.shape == (2, 3, 96, 96)
def test_unet_generator(): # color to color cfg = 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)) net = build_backbone(cfg) net.init_weights(pretrained=None) # cpu input_shape = (1, 3, 256, 256) img = _demo_inputs(input_shape) output = net(img) assert output.shape == (1, 3, 256, 256) # gpu if torch.cuda.is_available(): net = net.cuda() output = net(img.cuda()) assert output.shape == (1, 3, 256, 256) # gray to color cfg = dict(type='UnetGenerator', in_channels=1, 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)) net = build_backbone(cfg) net.init_weights(pretrained=None) # cpu input_shape = (1, 1, 256, 256) img = _demo_inputs(input_shape) output = net(img) assert output.shape == (1, 3, 256, 256) # gpu if torch.cuda.is_available(): net = net.cuda() output = net(img.cuda()) assert output.shape == (1, 3, 256, 256) # color to gray cfg = dict(type='UnetGenerator', in_channels=3, out_channels=1, num_down=8, base_channels=64, norm_cfg=dict(type='BN'), use_dropout=True, init_cfg=dict(type='normal', gain=0.02)) net = build_backbone(cfg) net.init_weights(pretrained=None) # cpu input_shape = (1, 3, 256, 256) img = _demo_inputs(input_shape) output = net(img) assert output.shape == (1, 1, 256, 256) # gpu if torch.cuda.is_available(): net = net.cuda() output = net(img.cuda()) assert output.shape == (1, 1, 256, 256) # pretrained should be str or None with pytest.raises(TypeError): net.init_weights(pretrained=[1]) # test norm_cfg assertions bad_cfg = copy.deepcopy(cfg) bad_cfg['norm_cfg'] = None with pytest.raises(AssertionError): _ = build_backbone(bad_cfg) bad_cfg['norm_cfg'] = dict(tp='BN') with pytest.raises(AssertionError): _ = build_backbone(bad_cfg)
def test_resnet_generator(): # color to color cfg = dict(type='ResnetGenerator', in_channels=3, out_channels=3, base_channels=64, norm_cfg=dict(type='IN'), use_dropout=False, num_blocks=9, padding_mode='reflect', init_cfg=dict(type='normal', gain=0.02)) net = build_backbone(cfg) net.init_weights(pretrained=None) # cpu input_shape = (1, 3, 256, 256) img = _demo_inputs(input_shape) output = net(img) assert output.shape == (1, 3, 256, 256) # gpu if torch.cuda.is_available(): net = net.cuda() output = net(img.cuda()) assert output.shape == (1, 3, 256, 256) # gray to color cfg = dict(type='ResnetGenerator', in_channels=1, out_channels=3, base_channels=64, norm_cfg=dict(type='IN'), use_dropout=False, num_blocks=9, padding_mode='reflect', init_cfg=dict(type='normal', gain=0.02)) net = build_backbone(cfg) net.init_weights(pretrained=None) # cpu input_shape = (1, 1, 256, 256) img = _demo_inputs(input_shape) output = net(img) assert output.shape == (1, 3, 256, 256) # gpu if torch.cuda.is_available(): net = net.cuda() output = net(img.cuda()) assert output.shape == (1, 3, 256, 256) # color to gray cfg = dict(type='ResnetGenerator', in_channels=3, out_channels=1, base_channels=64, norm_cfg=dict(type='IN'), use_dropout=False, num_blocks=9, padding_mode='reflect', init_cfg=dict(type='normal', gain=0.02)) net = build_backbone(cfg) net.init_weights(pretrained=None) # cpu input_shape = (1, 3, 256, 256) img = _demo_inputs(input_shape) output = net(img) assert output.shape == (1, 1, 256, 256) # gpu if torch.cuda.is_available(): net = net.cuda() output = net(img.cuda()) assert output.shape == (1, 1, 256, 256) # test num_blocks non-negative bad_cfg = copy.deepcopy(cfg) bad_cfg['num_blocks'] = -1 with pytest.raises(AssertionError): net = build_backbone(bad_cfg) # pretrained should be str or None with pytest.raises(TypeError): net.init_weights(pretrained=[1]) # test norm_cfg assertions bad_cfg = copy.deepcopy(cfg) bad_cfg['norm_cfg'] = None with pytest.raises(AssertionError): _ = build_backbone(bad_cfg) bad_cfg['norm_cfg'] = dict(tp='IN') with pytest.raises(AssertionError): _ = build_backbone(bad_cfg)
def test_cain_net(): model_cfg = dict(type='CAINNet') # build model model = build_backbone(model_cfg) # test attributes assert model.__class__.__name__ == 'CAINNet' # prepare data inputs0 = torch.rand(1, 2, 3, 5, 5) target0 = torch.rand(1, 3, 5, 5) inputs = torch.rand(1, 2, 3, 256, 248) target = torch.rand(1, 3, 256, 248) # test on cpu output = model(inputs) output = model(inputs, padding_flag=True) model(inputs0, padding_flag=True) assert torch.is_tensor(output) assert output.shape == target.shape with pytest.raises(AssertionError): output = model(inputs[:, :1]) with pytest.raises(OSError): model.init_weights('') with pytest.raises(TypeError): model.init_weights(1) model_cfg = dict(type='CAINNet', norm='in') model = build_backbone(model_cfg) model(inputs) model_cfg = dict(type='CAINNet', norm='bn') model = build_backbone(model_cfg) model(inputs) with pytest.raises(ValueError): model_cfg = dict(type='CAINNet', norm='lys') build_backbone(model_cfg) # test on gpu if torch.cuda.is_available(): model = model.cuda() inputs = inputs.cuda() target = target.cuda() output = model(inputs) output = model(inputs, True) assert torch.is_tensor(output) assert output.shape == target.shape inputs0 = inputs0.cuda() target0 = target0.cuda() model(inputs0, padding_flag=True) model_cfg = dict(type='CAINNet', norm='in') model = build_backbone(model_cfg).cuda() model(inputs) model_cfg = dict(type='CAINNet', norm='bn') model = build_backbone(model_cfg).cuda() model(inputs) with pytest.raises(ValueError): model_cfg = dict(type='CAINNet', norm='lys') build_backbone(model_cfg).cuda()