def test_lte(): model_cfg = dict(type='LTE', requires_grad=False, pixel_range=1., pretrained=None, load_pretrained_vgg=False) lte = build_component(model_cfg) assert lte.__class__.__name__ == 'LTE' x = torch.rand(2, 3, 64, 64) x_level3, x_level2, x_level1 = lte(x) assert x_level1.shape == (2, 64, 64, 64) assert x_level2.shape == (2, 128, 32, 32) assert x_level3.shape == (2, 256, 16, 16) lte.init_weights(None) with pytest.raises(IOError): model_cfg['pretrained'] = '' lte = build_component(model_cfg) x_level3, x_level2, x_level1 = lte(x) lte.init_weights('') with pytest.raises(TypeError): lte.init_weights(1)
def test_patch_discriminator(): # color, BN cfg = dict(type='PatchDiscriminator', in_channels=3, base_channels=64, num_conv=3, norm_cfg=dict(type='BN'), init_cfg=dict(type='normal', gain=0.02)) net = build_component(cfg) net.init_weights(pretrained=None) # cpu input_shape = (1, 3, 64, 64) img = _demo_inputs(input_shape) output = net(img) assert output.shape == (1, 1, 6, 6) # gpu if torch.cuda.is_available(): net.init_weights(pretrained=None) net = net.cuda() output = net(img.cuda()) assert output.shape == (1, 1, 6, 6) # pretrained should be str or None with pytest.raises(TypeError): net.init_weights(pretrained=[1]) # gray, IN cfg = dict(type='PatchDiscriminator', in_channels=1, base_channels=64, num_conv=3, norm_cfg=dict(type='IN'), init_cfg=dict(type='normal', gain=0.02)) net = build_component(cfg) net.init_weights(pretrained=None) # cpu input_shape = (1, 1, 64, 64) img = _demo_inputs(input_shape) output = net(img) assert output.shape == (1, 1, 6, 6) # gpu if torch.cuda.is_available(): net.init_weights(pretrained=None) net = net.cuda() output = net(img.cuda()) assert output.shape == (1, 1, 6, 6) # 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_component(bad_cfg) bad_cfg['norm_cfg'] = dict(tp='BN') with pytest.raises(AssertionError): _ = build_component(bad_cfg)
def test_smpatch_discriminator(): # color, BN cfg = dict(type='SoftMaskPatchDiscriminator', in_channels=3, base_channels=64, num_conv=3, with_spectral_norm=True) net = build_component(cfg) net.init_weights(pretrained=None) # cpu input_shape = (1, 3, 64, 64) img = _demo_inputs(input_shape) output = net(img) assert output.shape == (1, 1, 6, 6) # gpu if torch.cuda.is_available(): net.init_weights(pretrained=None) net = net.cuda() output = net(img.cuda()) assert output.shape == (1, 1, 6, 6) # pretrained should be str or None with pytest.raises(TypeError): net.init_weights(pretrained=[1]) # gray, IN cfg = dict(type='SoftMaskPatchDiscriminator', in_channels=1, base_channels=64, num_conv=3, with_spectral_norm=True) net = build_component(cfg) net.init_weights(pretrained=None) # cpu input_shape = (1, 1, 64, 64) img = _demo_inputs(input_shape) output = net(img) assert output.shape == (1, 1, 6, 6) # gpu if torch.cuda.is_available(): net.init_weights(pretrained=None) net = net.cuda() output = net(img.cuda()) assert output.shape == (1, 1, 6, 6) # pretrained should be str or None with pytest.raises(TypeError): net.init_weights(pretrained=[1])
def test_gl_discs(): global_disc_cfg = dict(in_channels=3, max_channels=512, fc_in_channels=512 * 4 * 4, fc_out_channels=1024, num_convs=6, norm_cfg=dict(type='BN')) local_disc_cfg = dict(in_channels=3, max_channels=512, fc_in_channels=512 * 4 * 4, fc_out_channels=1024, num_convs=5, norm_cfg=dict(type='BN')) gl_disc_cfg = dict(type='GLDiscs', global_disc_cfg=global_disc_cfg, local_disc_cfg=local_disc_cfg) gl_discs = build_component(gl_disc_cfg) gl_discs.init_weights() input_g = torch.randn(1, 3, 256, 256) input_l = torch.randn(1, 3, 128, 128) output = gl_discs((input_g, input_l)) assert output.shape == (1, 1) with pytest.raises(TypeError): gl_discs.init_weights(pretrained=dict(igccc=777)) if torch.cuda.is_available(): gl_discs = gl_discs.cuda() input_g = torch.randn(1, 3, 256, 256).cuda() input_l = torch.randn(1, 3, 128, 128).cuda() output = gl_discs((input_g, input_l)) assert output.shape == (1, 1)
def test_feedback_hour_glass(): model_cfg = dict(type='FeedbackHourglass', mid_channels=16, num_keypoints=20) fhg = build_component(model_cfg) assert fhg.__class__.__name__ == 'FeedbackHourglass' x = torch.rand(2, 3, 64, 64) heatmap, last_hidden = fhg.forward(x) assert heatmap.shape == (2, 20, 16, 16) assert last_hidden.shape == (2, 16, 16, 16) heatmap, last_hidden = fhg.forward(x, last_hidden) assert heatmap.shape == (2, 20, 16, 16) assert last_hidden.shape == (2, 16, 16, 16)
def test_search_transformer(): model_cfg = dict(type='SearchTransformer') model = build_component(model_cfg) lr_pad_level3 = torch.randn((2, 32, 32, 32)) ref_pad_level3 = torch.randn((2, 32, 32, 32)) ref_level3 = torch.randn((2, 32, 32, 32)) ref_level2 = torch.randn((2, 16, 64, 64)) ref_level1 = torch.randn((2, 8, 128, 128)) s, t_level3, t_level2, t_level1 = model(lr_pad_level3, ref_pad_level3, ref_level1, ref_level2, ref_level3) assert s.shape == (2, 1, 32, 32) assert t_level3.shape == (2, 32, 32, 32) assert t_level2.shape == (2, 16, 64, 64) assert t_level1.shape == (2, 8, 128, 128)
def test_ttsr_dict(): cfg = dict(type='TTSRDiscriminator', in_channels=3, in_size=160) net = build_component(cfg) net.init_weights(pretrained=None) # cpu inputs = torch.rand((2, 3, 160, 160)) output = net(inputs) assert output.shape == (2, 1) # gpu if torch.cuda.is_available(): net.init_weights(pretrained=None) net = net.cuda() output = net(inputs.cuda()) assert output.shape == (2, 1) # pretrained should be str or None with pytest.raises(TypeError): net.init_weights(pretrained=[1])
def __init__(self, global_disc_cfg, local_disc_cfg): super(DeepFillv1Discriminators, self).__init__() self.global_disc = build_component(global_disc_cfg) self.local_disc = build_component(local_disc_cfg)
def __init__(self, global_disc_cfg, local_disc_cfg): super().__init__() self.global_disc = build_component(global_disc_cfg) self.local_disc = build_component(local_disc_cfg)