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)
