def test_biggan_deep_discriminator_cuda(self):
# test default settings
d = build_module(self.default_config).cuda()
assert isinstance(d, BigGANDeepDiscriminator)
y = d(self.x.cuda(), self.label.cuda())
assert y.shape == (2, 1)
# test different init types
cfg = deepcopy(self.default_config)
cfg.update(dict(init_type='N02'))
d = build_module(cfg).cuda()
y = d(self.x.cuda(), self.label.cuda())
assert y.shape == (2, 1)
cfg = deepcopy(self.default_config)
cfg.update(dict(init_type='xavier'))
d = build_module(cfg).cuda()
y = d(self.x.cuda(), self.label.cuda())
assert y.shape == (2, 1)
# test different num_classes
cfg = deepcopy(self.default_config)
cfg.update(dict(num_classes=0))
d = build_module(cfg).cuda()
y = d(self.x.cuda(), None)
assert y.shape == (2, 1)
# test with `with_spectral_norm=False`
cfg = deepcopy(self.default_config)
cfg.update(dict(with_spectral_norm=False))
d = build_module(cfg).cuda()
y = d(self.x.cuda(), self.label.cuda())
assert y.shape == (2, 1)
def test_biggan_condition_bn_cuda(self):
# test default setting
module = build_module(self.default_cfg).cuda()
assert isinstance(module, BigGANConditionBN)
out = module(self.x.cuda(), self.y.cuda())
assert out.shape == (2, 64, 32, 32)
# test input_is_label
cfg = deepcopy(self.default_cfg)
cfg.update(dict(input_is_label=True))
module = build_module(cfg).cuda()
out = module(self.x.cuda(), self.label.cuda())
assert out.shape == (2, 64, 32, 32)
# test torch-sn
cfg = deepcopy(self.default_cfg)
cfg.update(dict(sn_style='torch'))
module = build_module(cfg).cuda()
out = module(self.x.cuda(), self.y.cuda())
assert out.shape == (2, 64, 32, 32)
# test not-implemented sn-style
with pytest.raises(NotImplementedError):
cfg = deepcopy(self.default_cfg)
cfg.update(dict(sn_style='tero'))
module = build_module(cfg).cuda()
out = module(self.x.cuda(), self.y.cuda())
assert out.shape == (2, 64, 32, 32)
def test_wgangp_discriminator_cuda(self):
# test default setting with builder
d = build_module(self.default_config).cuda()
assert isinstance(d, WGANGPDiscriminator)
score = d(self.x.cuda())
assert score.shape == (2, 1)
# test different in_scale
config = dict(type='WGANGPDiscriminator', in_channel=3, in_scale=64)
d = build_module(config).cuda()
assert isinstance(d, WGANGPDiscriminator)
x = torch.randn((2, 3, 64, 64))
score = d(x.cuda())
assert score.shape == (2, 1)
# test different conv config
config = dict(type='WGANGPDiscriminator',
in_channel=3,
in_scale=128,
conv_module_cfg=self.conv_ln_module_config)
d = build_module(config).cuda()
assert isinstance(d, WGANGPDiscriminator)
score = d(self.x.cuda())
assert score.shape == (2, 1)
config = dict(type='WGANGPDiscriminator',
in_channel=3,
in_scale=128,
conv_module_cfg=self.conv_gn_module_config)
d = build_module(config).cuda()
assert isinstance(d, WGANGPDiscriminator)
score = d(self.x.cuda())
assert score.shape == (2, 1)
def test_wgangp_generator_cuda(self):
# test default setting with builder
g = build_module(self.default_config).cuda()
assert isinstance(g, WGANGPGenerator)
x = g(None, num_batches=3)
assert x.shape == (3, 3, 128, 128)
# test different out_scale
config = dict(type='WGANGPGenerator', noise_size=128, out_scale=64)
g = build_module(config).cuda()
assert isinstance(g, WGANGPGenerator)
x = g(None, num_batches=3)
assert x.shape == (3, 3, 64, 64)
# test different conv config
config = dict(
type='WGANGPGenerator',
noise_size=128,
out_scale=128,
conv_module_cfg=dict(
conv_cfg=None,
kernel_size=3,
stride=1,
padding=1,
bias=True,
act_cfg=dict(type='LeakyReLU', negative_slope=0.2),
norm_cfg=dict(type='BN'),
order=('conv', 'norm', 'act')))
g = build_module(config).cuda()
assert isinstance(g, WGANGPGenerator)
x = g(None, num_batches=3)
assert x.shape == (3, 3, 128, 128)
def test_snganDiscHeadResBlock_cuda(self):
# test default config
block = build_module(self.default_config).cuda()
out = block(self.input.cuda())
assert out.shape == (2, 16, 5, 5)
# test init cfg + w or w/o downsample
for init_method in [
'studio', 'biggan', 'sagan', 'sngan', 'sngan-proj', 'gan-proj'
]:
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type=init_method)
block = build_module(config)
out = block(self.input)
assert out.shape == (2, 16, 5, 5)
# test conv_cfg
config = deepcopy(self.default_config)
config['conv_cfg'] = dict(kernel_size=1,
stride=1,
padding=0,
act_cfg=None)
block = build_module(config).cuda()
out = block(self.input.cuda())
assert out.shape == (2, 16, 5, 5)
def test_biggan_deep_disc_res_block_cuda(self):
# test default setting
module = build_module(self.default_cfg).cuda()
assert isinstance(module, BigGANDeepDiscResBlock)
out = module(self.x.cuda())
assert out.shape == (2, 64, 8, 8)
# test with_downsample
cfg = deepcopy(self.default_cfg)
cfg.update(dict(with_downsample=False))
module = build_module(cfg).cuda()
out = module(self.x.cuda())
assert out.shape == (2, 64, 16, 16)
# test different channel_ratio
cfg = deepcopy(self.default_cfg)
cfg.update(dict(channel_ratio=8))
module = build_module(cfg)
out = module(self.x)
assert out.shape == (2, 64, 8, 8)
# test torch-sn
cfg = deepcopy(self.default_cfg)
cfg.update(dict(sn_style='torch'))
module = build_module(cfg).cuda()
out = module(self.x.cuda())
assert out.shape == (2, 64, 8, 8)
def test_biggan_deep_gen_res_block_cuda(self):
# test default setting
module = build_module(self.default_cfg).cuda()
assert isinstance(module, BigGANDeepGenResBlock)
out = module(self.x.cuda(), self.y.cuda())
assert out.shape == (2, 16, 16, 16)
# test without upsample
cfg = deepcopy(self.default_cfg)
cfg.update(dict(upsample_cfg=None))
module = build_module(cfg).cuda()
out = module(self.x.cuda(), self.y.cuda())
assert out.shape == (2, 16, 8, 8)
# test input_is_label
cfg = deepcopy(self.default_cfg)
cfg.update(dict(input_is_label=True))
module = build_module(cfg).cuda()
out = module(self.x.cuda(), self.label.cuda())
assert out.shape == (2, 16, 16, 16)
# test torch-sn
cfg = deepcopy(self.default_cfg)
cfg.update(dict(sn_style='torch'))
module = build_module(cfg).cuda()
out = module(self.x.cuda(), self.y.cuda())
assert out.shape == (2, 16, 16, 16)
def test_lsgan_discriminator_cuda(self):
# test default setting with builder
d = build_module(self.default_config).cuda()
assert isinstance(d, LSGANDiscriminator)
score = d(self.x.cuda())
assert score.shape == (2, 1)
# test different input_scale
config = dict(type='LSGANDiscriminator', in_channels=3, input_scale=64)
d = build_module(config).cuda()
assert isinstance(d, LSGANDiscriminator)
x = torch.randn((2, 3, 64, 64))
score = d(x.cuda())
assert score.shape == (2, 1)
# test different config
config = dict(
type='LSGANDiscriminator',
in_channels=3,
input_scale=64,
out_act_cfg=dict(type='Sigmoid'))
d = build_module(config).cuda()
assert isinstance(d, LSGANDiscriminator)
x = torch.randn((2, 3, 64, 64))
score = d(x.cuda())
assert score.shape == (2, 1)
def test_snganDiscResBlock(self):
# test default config
block = build_module(self.default_config)
out = block(self.input)
assert out.shape == (2, 16, 5, 5)
# test conv_cfg
config = deepcopy(self.default_config)
config['conv_cfg'] = dict(
kernel_size=1, stride=1, padding=0, act_cfg=None)
block = build_module(config)
out = block(self.input)
assert out.shape == (2, 16, 5, 5)
# test w/o learnabel shortcut + w/o downsample
config = deepcopy(self.default_config)
config['downsample'] = False
config['out_channels'] = 8
block = build_module(config)
out = block(self.input)
assert out.shape == (2, 8, 10, 10)
# test init_cft --> raise error
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='wgan-gp')
with pytest.raises(NotImplementedError):
block = build_module(config)
def test_snganDiscHeadResBlock(self):
# test default config
block = build_module(self.default_config)
out = block(self.input)
assert out.shape == (2, 16, 5, 5)
# test conv_cfg
config = deepcopy(self.default_config)
config['conv_cfg'] = dict(kernel_size=1,
stride=1,
padding=0,
act_cfg=None)
block = build_module(config)
out = block(self.input)
assert out.shape == (2, 16, 5, 5)
# test init cfg + w or w/o downsample
for init_method in [
'studio', 'biggan', 'sagan', 'sngan', 'sngan-proj', 'gan-proj'
]:
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type=init_method)
block = build_module(config)
out = block(self.input)
assert out.shape == (2, 16, 5, 5)
# test init_cft --> raise error
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='wgan-gp')
with pytest.raises(NotImplementedError):
block = build_module(config)
def test_biggan_condition_bn_cuda(self):
# test default setting
module = build_module(self.default_cfg).cuda()
assert isinstance(module, BigGANConditionBN)
out = module(self.x.cuda(), self.y.cuda())
assert out.shape == (2, 64, 32, 32)
# test input_is_label
cfg = deepcopy(self.default_cfg)
cfg.update(dict(input_is_label=True))
module = build_module(cfg).cuda()
out = module(self.x.cuda(), self.label.cuda())
assert out.shape == (2, 64, 32, 32)
def test_snganDiscHeadResBlock_cuda(self):
# test default config
block = build_module(self.default_config).cuda()
out = block(self.input.cuda())
assert out.shape == (2, 16, 5, 5)
# test conv_cfg
config = deepcopy(self.default_config)
config['conv_cfg'] = dict(
kernel_size=1, stride=1, padding=0, act_cfg=None)
block = build_module(config).cuda()
out = block(self.input.cuda())
assert out.shape == (2, 16, 5, 5)
def test_biggan_disc_res_block_cuda(self):
# test default setting
module = build_module(self.default_cfg).cuda()
assert isinstance(module, BigGANDiscResBlock)
out = module(self.x.cuda())
assert out.shape == (2, 64, 8, 8)
# test with_downsample
cfg = deepcopy(self.default_cfg)
cfg.update(dict(with_downsample=False))
module = build_module(cfg).cuda()
out = module(self.x.cuda())
assert out.shape == (2, 64, 16, 16)
def test_self_attention_block_cuda(self):
# test default setting
module = build_module(self.default_cfg).cuda()
assert isinstance(module, SelfAttentionBlock)
out = module(self.x.cuda())
assert out.shape == (2, 16, 8, 8)
# test different in_channels
cfg = deepcopy(self.default_cfg)
cfg.update(dict(in_channels=10))
module = build_module(cfg).cuda()
x = torch.randn(2, 10, 8, 8).cuda()
out = module(x)
assert out.shape == (2, 10, 8, 8)
def test_sngan_proj_discriminator(self):
# test default setting with builder
d = build_module(self.default_config)
assert isinstance(d, SAGANDiscriminator)
score = d(self.x, self.label)
assert score.shape == (2, 1)
# test different input_scale
config = deepcopy(self.default_config)
config['input_scale'] = 64
d = build_module(config)
assert isinstance(d, SAGANDiscriminator)
x = torch.randn((2, 3, 64, 64))
score = d(x, self.label)
assert score.shape == (2, 1)
# test num_classes == 0 (w/o proj_y)
config = deepcopy(self.default_config)
config['num_classes'] = 0
d = build_module(config)
assert isinstance(d, SAGANDiscriminator)
score = d(self.x, self.label)
assert score.shape == (2, 1)
# test different base_channels
config = deepcopy(self.default_config)
config['base_channels'] = 128
d = build_module(config)
assert isinstance(d, SAGANDiscriminator)
score = d(self.x, self.label)
assert score.shape == (2, 1)
# test different attention_after_nth_block
config = deepcopy(self.default_config)
config['attention_after_nth_block'] = [1, 2]
d = build_module(config)
assert isinstance(d, SAGANDiscriminator)
score = d(self.x, self.label)
assert score.shape == (2, 1)
# wrong type of attention_after_nth_block --> wrong type
config = deepcopy(self.default_config)
config['attention_after_nth_block'] = '1'
with pytest.raises(ValueError):
d = build_module(config)
# wrong type of attention_after_nth_block --> wrong type of list
config = deepcopy(self.default_config)
config['attention_after_nth_block'] = ['1', '2']
with pytest.raises(ValueError):
d = build_module(config)
# test init_cfg --> SAGAN
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='sagan')
d = build_module(config)
assert isinstance(d, SAGANDiscriminator)
score = d(self.x, self.label)
assert score.shape == (2, 1)
def test_conditionalNorm(self):
# test build from default config
norm = build_module(self.default_config)
out = norm(self.input, self.label)
assert out.shape == (2, 4, 4, 4)
# test w/o use_cbn
config = deepcopy(self.default_config)
config['use_cbn'] = False
norm = build_module(config)
out = norm(self.input)
assert out.shape == (2, 4, 4, 4)
# test num_class < 0 and cbn = False
config = deepcopy(self.default_config)
config['num_classes'] = 0
config['use_cbn'] = False
norm = build_module(config)
out = norm(self.input)
assert out.shape == (2, 4, 4, 4)
# test num_classes <= 0 and cbn = True
config = deepcopy(self.default_config)
config['num_classes'] = 0
with pytest.raises(ValueError):
norm = build_module(config)
# test IN
config = deepcopy(self.default_config)
config['norm_cfg'] = dict(type='IN')
norm = build_module(config)
out = norm(self.input, self.label)
assert out.shape == (2, 4, 4, 4)
# test SyncBN
# config = deepcopy(self.default_config)
# config['norm_cfg'] = dict(type='SyncBN')
# norm = build_module(config)
# out = norm(self.input, self.label)
# assert out.shape == (2, 4, 4, 4)
# test unknown norm type
config = deepcopy(self.default_config)
config['norm_cfg'] = dict(type='GN')
with pytest.raises(ValueError):
norm = build_module(config)
# test init_cfg
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='sngan')
norm = build_module(config)
out = norm(self.input, self.label)
assert out.shape == (2, 4, 4, 4)
# test init_cft --> raise error
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='wgan-gp')
with pytest.raises(NotImplementedError):
norm = build_module(config)
def test_snganDiscResBlock(self):
# test default config
block = build_module(self.default_config)
out = block(self.input)
assert out.shape == (2, 16, 5, 5)
# test conv_cfg
config = deepcopy(self.default_config)
config['conv_cfg'] = dict(kernel_size=1,
stride=1,
padding=0,
act_cfg=None)
block = build_module(config)
out = block(self.input)
assert out.shape == (2, 16, 5, 5)
# test w/o learnabel shortcut + w/o downsample
config = deepcopy(self.default_config)
config['downsample'] = False
config['out_channels'] = 8
block = build_module(config)
out = block(self.input)
assert out.shape == (2, 8, 10, 10)
# test init cfg + w or w/o downsample
for init_method in [
'studio', 'biggan', 'sagan', 'sngan', 'sngan-proj', 'gan-proj'
]:
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type=init_method)
config['out_channels'] = 8
for downsample in [True, False]:
config['downsample'] = downsample
block = build_module(config)
out = block(self.input)
if downsample:
assert out.shape == (2, 8, 5, 5)
else:
assert out.shape == (2, 8, 10, 10)
# test init_cft --> raise error
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='wgan-gp')
with pytest.raises(NotImplementedError):
block = build_module(config)
def test_sagan_generator_cuda(self):
# test default setting with builder
g = build_module(self.default_config).cuda()
assert isinstance(g, SAGANGenerator)
x = g(None, num_batches=2)
assert x.shape == (2, 3, 32, 32)
# test return noise
x = g(None, num_batches=2, return_noise=True)
assert x['fake_img'].shape == (2, 3, 32, 32)
assert x['noise_batch'].shape == (2, 128)
assert x['label'].shape == (2, )
x = g(self.noise.cuda(), label=self.label.cuda(), return_noise=True)
assert x['noise_batch'].shape == (2, 128)
assert x['label'].shape == (2, )
x = g(torch.randn, num_batches=2, return_noise=True)
assert x['noise_batch'].shape == (2, 128)
assert x['label'].shape == (2, )
# test different output_scale
config = deepcopy(self.default_config)
config['output_scale'] = 64
g = build_module(config).cuda()
assert isinstance(g, SAGANGenerator)
x = g(None, num_batches=2)
assert x.shape == (2, 3, 64, 64)
# test different attention_after_nth_block
config = deepcopy(self.default_config)
config['attention_after_nth_block'] = [1, 2]
g = build_module(config).cuda()
assert isinstance(g, SAGANGenerator)
x = g(None, num_batches=2)
assert x.shape == (2, 3, 32, 32)
# test init_cfg --> SAGAN
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='sagan')
g = build_module(config).cuda()
assert isinstance(g, SAGANGenerator)
x = g(None, num_batches=2)
assert x.shape == (2, 3, 32, 32)
def test_snganDiscHeadResBlock(self):
# test default config
block = build_module(self.default_config)
out = block(self.input)
assert out.shape == (2, 16, 5, 5)
# test conv_cfg
config = deepcopy(self.default_config)
config['conv_cfg'] = dict(
kernel_size=1, stride=1, padding=0, act_cfg=None)
block = build_module(config)
out = block(self.input)
assert out.shape == (2, 16, 5, 5)
# test init_cft --> raise error
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='wgan-gp')
with pytest.raises(NotImplementedError):
block = build_module(config)
def test_biggan_gen_res_block(self):
# test default setting
module = build_module(self.default_cfg)
assert isinstance(module, BigGANGenResBlock)
out = module(self.x, self.y)
assert out.shape == (2, 16, 16, 16)
# test without upsample
cfg = deepcopy(self.default_cfg)
cfg.update(dict(upsample_cfg=None))
module = build_module(cfg)
out = module(self.x, self.y)
assert out.shape == (2, 16, 8, 8)
# test input_is_label
cfg = deepcopy(self.default_cfg)
cfg.update(dict(input_is_label=True))
module = build_module(cfg)
out = module(self.x, self.label)
assert out.shape == (2, 16, 16, 16)
def test_lsgan_generator_cuda(self):
# test default setting with builder
g = build_module(self.default_config).cuda()
assert isinstance(g, LSGANGenerator)
x = g(None, num_batches=3)
assert x.shape == (3, 3, 128, 128)
x = g(None, num_batches=3, return_noise=True)
assert x['noise_batch'].shape == (3, 128)
x = g(self.noise.cuda(), return_noise=True)
assert x['noise_batch'].shape == (3, 128)
x = g(torch.randn, num_batches=3, return_noise=True)
assert x['noise_batch'].shape == (3, 128)
# test different output_scale
config = dict(type='LSGANGenerator', noise_size=128, output_scale=64)
g = build_module(config).cuda()
assert isinstance(g, LSGANGenerator)
x = g(None, num_batches=3)
assert x.shape == (3, 3, 64, 64)
def test_snganDiscResBlock_cuda(self):
# test default config
block = build_module(self.default_config).cuda()
out = block(self.input.cuda())
assert out.shape == (2, 16, 5, 5)
# test conv_cfg
config = deepcopy(self.default_config)
config['conv_cfg'] = dict(
kernel_size=1, stride=1, padding=0, act_cfg=None)
block = build_module(config).cuda()
out = block(self.input.cuda())
assert out.shape == (2, 16, 5, 5)
# test w/o learnabel shortcut + w/o downsample
config = deepcopy(self.default_config)
config['downsample'] = False
config['out_channels'] = 8
block = build_module(config).cuda()
out = block(self.input.cuda())
assert out.shape == (2, 8, 10, 10)
def test_dcgan_discriminator(self):
# test default setting with builder
d = build_module(self.default_config)
pred = d(self.input_tensor)
assert pred.shape == (2, 5)
assert d.num_downsamples == 3
assert len(d.downsamples) == 3
assert not d.downsamples[0].with_norm
assert not d.output_layer.with_norm
assert not d.output_layer.with_activation
assert isinstance(d.downsamples[1].activate, torch.nn.LeakyReLU)
assert isinstance(d.downsamples[1].norm, torch.nn.BatchNorm2d)
# sanity check for args with cpu model
d = DCGANDiscriminator(input_scale=64, output_scale=8, out_channels=2)
assert d.input_scale == 64 and d.output_scale == 8
assert d.num_downsamples == 3
assert d.out_channels == 2
pred = d(torch.randn((1, 3, 64, 64)))
assert pred.shape == (1, 50)
with pytest.raises(TypeError):
_ = DCGANDiscriminator(32, 4, 2, pretrained=dict())
# check for cuda
if not torch.cuda.is_available():
return
# test default setting with builder on GPU
d = build_module(self.default_config).cuda()
pred = d(self.input_tensor.cuda())
assert pred.shape == (2, 5)
assert d.num_downsamples == 3
assert len(d.downsamples) == 3
assert not d.downsamples[0].with_norm
assert not d.output_layer.with_norm
assert not d.output_layer.with_activation
assert isinstance(d.downsamples[1].activate, torch.nn.LeakyReLU)
def test_snganGenResBlock_cuda(self):
# test default config
block = build_module(self.default_config).cuda()
out = block(self.input.cuda(), self.label.cuda())
assert out.shape == (2, 16, 10, 10)
# test no upsample config and no learnable sc
config = deepcopy(self.default_config)
config['upsample'] = False
block = build_module(config).cuda()
out = block(self.input.cuda(), self.label.cuda())
assert out.shape == (2, 16, 5, 5)
# test learnable shortcut + w/o upsample
config = deepcopy(self.default_config)
config['out_channels'] = 32
config['upsample'] = False
block = build_module(config).cuda()
out = block(self.input.cuda(), self.label.cuda())
assert out.shape == (2, 32, 5, 5)
# test init_cfg + w/o learnable shortcut
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='sngan')
config['upsample'] = False
block = build_module(config).cuda()
out = block(self.input.cuda(), self.label.cuda())
assert out.shape == (2, 16, 5, 5)
# test conv_cfg
config = deepcopy(self.default_config)
config['conv_cfg'] = dict(
kernel_size=1, stride=1, padding=0, act_cfg=None)
block = build_module(config).cuda()
out = block(self.input.cuda(), self.label.cuda())
assert out.shape == (2, 16, 10, 10)
def test_lsgan_discriminator_cuda(self):
# test default setting with builder
d = build_module(self.default_config).cuda()
assert isinstance(d, SAGANDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
# test different input_scale
config = deepcopy(self.default_config)
config['input_scale'] = 64
d = build_module(config).cuda()
assert isinstance(d, SAGANDiscriminator)
x = torch.randn((2, 3, 64, 64)).cuda()
score = d(x, self.label.cuda())
assert score.shape == (2, 1)
# test num_classes == 0 (w/o proj_y)
config = deepcopy(self.default_config)
config['num_classes'] = 0
d = build_module(config).cuda()
assert isinstance(d, SAGANDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
# test different base_channels
config = deepcopy(self.default_config)
config['base_channels'] = 128
d = build_module(config).cuda()
assert isinstance(d, SAGANDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
# test different attention_after_nth_block
config = deepcopy(self.default_config)
config['attention_after_nth_block'] = [1, 2]
d = build_module(config).cuda()
assert isinstance(d, SAGANDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
# test init_cfg --> SAGAN
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='sagan')
d = build_module(config).cuda()
assert isinstance(d, SAGANDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
def test_snganGenResBlock(self):
# test default config
block = build_module(self.default_config)
out = block(self.input, self.label)
assert out.shape == (2, 16, 10, 10)
# test no upsample config and no learnable sc
config = deepcopy(self.default_config)
config['upsample'] = False
block = build_module(config)
out = block(self.input, self.label)
assert out.shape == (2, 16, 5, 5)
# test learnable shortcut + w/o upsample
config = deepcopy(self.default_config)
config['out_channels'] = 32
config['upsample'] = False
block = build_module(config)
out = block(self.input, self.label)
assert out.shape == (2, 32, 5, 5)
# test init_cfg + w/o learnable shortcut
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='sngan')
config['upsample'] = False
block = build_module(config)
out = block(self.input, self.label)
assert out.shape == (2, 16, 5, 5)
# test init_cft --> raise error
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='wgan-gp')
with pytest.raises(NotImplementedError):
block = build_module(config)
# test conv_cfg
config = deepcopy(self.default_config)
config['conv_cfg'] = dict(
kernel_size=1, stride=1, padding=0, act_cfg=None)
block = build_module(config)
out = block(self.input, self.label)
assert out.shape == (2, 16, 10, 10)
def test_conditionalNorm_cuda(self):
# test build from default config
norm = build_module(self.default_config).cuda()
out = norm(self.input.cuda(), self.label.cuda())
assert out.shape == (2, 4, 4, 4)
# test w/o use_cbn
config = deepcopy(self.default_config)
config['use_cbn'] = False
norm = build_module(config).cuda()
out = norm(self.input.cuda())
assert out.shape == (2, 4, 4, 4)
# test num_class < 0 and cbn = False
config = deepcopy(self.default_config)
config['num_classes'] = 0
config['use_cbn'] = False
norm = build_module(config).cuda()
out = norm(self.input.cuda())
assert out.shape == (2, 4, 4, 4)
# test num_classes <= 0 and cbn = True
config = deepcopy(self.default_config)
config['num_classes'] = 0
with pytest.raises(ValueError):
norm = build_module(config)
# test IN
config = deepcopy(self.default_config)
config['norm_cfg'] = dict(type='IN')
norm = build_module(config).cuda()
out = norm(self.input.cuda(), self.label.cuda())
assert out.shape == (2, 4, 4, 4)
# test sn_style == ajbrock
config = deepcopy(self.default_config)
config['with_spectral_norm'] = True
config['sn_style'] = 'ajbrock'
norm = build_module(config)
out = norm(self.input, self.label)
for buffer in ['u0', 'sv0']:
assert hasattr(norm.weight_embedding, buffer)
assert hasattr(norm.bias_embedding, buffer)
assert out.shape == (2, 4, 4, 4)
# test sn_style == torch
config = deepcopy(self.default_config)
config['with_spectral_norm'] = True
config['sn_style'] = 'torch'
norm = build_module(config)
out = norm(self.input, self.label)
for buffer in ['weight_u', 'weight_v', 'weight_orig']:
assert hasattr(norm.weight_embedding, buffer)
assert hasattr(norm.bias_embedding, buffer)
assert out.shape == (2, 4, 4, 4)
# test sn_style --> raise error
config = deepcopy(self.default_config)
config['with_spectral_norm'] = True
config['sn_style'] = 'studio'
with pytest.raises(NotImplementedError):
norm = build_module(config)
# test SyncBN
# config = deepcopy(self.default_config)
# config['norm_cfg'] = dict(type='SyncBN')
# norm = build_module(config)
# out = norm(self.input, self.label)
# assert out.shape == (2, 4, 4, 4)
# test unknown norm type
config = deepcopy(self.default_config)
config['norm_cfg'] = dict(type='GN')
with pytest.raises(ValueError):
norm = build_module(config)
# test init_cfg
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='sngan')
norm = build_module(config).cuda()
out = norm(self.input.cuda(), self.label.cuda())
assert out.shape == (2, 4, 4, 4)
def test_dcgan_generator(self):
# test default setting with builder
g = build_module(self.default_config)
assert isinstance(g, DCGANGenerator)
assert g.num_upsamples == 2
assert not g.output_layer.with_norm
assert len(g.upsampling) == 1
assert g.upsampling[0].with_norm
assert g.noise2feat.with_norm
assert isinstance(g.output_layer.activate, torch.nn.Tanh)
# check forward function
img = g(self.noise)
assert img.shape == (2, 3, 16, 16)
img = g(self.noise[:, :, None, None])
assert img.shape == (2, 3, 16, 16)
img = g(torch.randn, num_batches=2)
assert img.shape == (2, 3, 16, 16)
img = g(None, num_batches=2)
assert img.shape == (2, 3, 16, 16)
with pytest.raises(ValueError):
_ = g(torch.randn((1, 100, 3)))
with pytest.raises(AssertionError):
_ = g(torch.randn)
with pytest.raises(AssertionError):
_ = g(None)
with pytest.raises(AssertionError):
_ = g(torch.randn(2, 10))
results = g(self.noise, return_noise=True)
assert results['noise_batch'].shape == (2, 100, 1, 1)
# sanity check for args with cpu model
g = DCGANGenerator(32, base_channels=64)
img = g(self.noise)
assert img.shape == (2, 3, 32, 32)
assert g.base_channels == 64
g = DCGANGenerator(16, out_channels=1, base_channels=32)
img = g(self.noise)
assert img.shape == (2, 1, 16, 16)
g = DCGANGenerator(16, noise_size=10, base_channels=32)
with pytest.raises(AssertionError):
_ = g(self.noise)
img = g(torch.randn(2, 10))
assert img.shape == (2, 3, 16, 16)
g = DCGANGenerator(16,
default_act_cfg=dict(type='LeakyReLU'),
base_channels=32)
assert isinstance(g.noise2feat.activate, torch.nn.LeakyReLU)
assert isinstance(g.upsampling[0].activate, torch.nn.LeakyReLU)
assert isinstance(g.output_layer.activate, torch.nn.Tanh)
with pytest.raises(TypeError):
_ = DCGANGenerator(16,
noise_size=10,
base_channels=32,
pretrained=dict())
# check for cuda
if not torch.cuda.is_available():
return
g = build_module(self.default_config).cuda()
assert isinstance(g, DCGANGenerator)
assert g.num_upsamples == 2
assert not g.output_layer.with_norm
assert len(g.upsampling) == 1
assert g.upsampling[0].with_norm
# check forward function
img = g(self.noise)
assert img.shape == (2, 3, 16, 16)
img = g(self.noise[:, :, None, None])
assert img.shape == (2, 3, 16, 16)
img = g(torch.randn, num_batches=2)
assert img.shape == (2, 3, 16, 16)
img = g(None, num_batches=2)
assert img.shape == (2, 3, 16, 16)
with pytest.raises(ValueError):
_ = g(torch.randn((1, 100, 3)))
with pytest.raises(AssertionError):
_ = g(torch.randn)
with pytest.raises(AssertionError):
_ = g(None)
results = g(self.noise, return_noise=True)
assert results['noise_batch'].shape == (2, 100, 1, 1)
def test_biggan_deep_generator_cuda(self):
# test default setting with builder
g = build_module(self.default_config).cuda()
assert isinstance(g, BigGANDeepGenerator)
res = g(self.noise.cuda(), self.label.cuda())
assert res.shape == (3, 3, 128, 128)
# test 'return_noise'
res = g(self.noise.cuda(), self.label.cuda(), return_noise=True)
assert res['fake_img'].shape == (3, 3, 128, 128)
assert res['noise_batch'].shape == (3, 120)
assert res['label'].shape == (3, )
res = g(None, None, num_batches=3, return_noise=True)
assert res['fake_img'].shape == (3, 3, 128, 128)
assert res['noise_batch'].shape == (3, 120)
assert res['label'].shape == (3, )
# test callable
noise = torch.randn
label = partial(torch.randint, 0, 1000)
res = g(noise, label, num_batches=2)
assert res.shape == (2, 3, 128, 128)
# test different output scale
cfg = deepcopy(self.default_config)
cfg.update(dict(output_scale=256))
g = build_module(cfg).cuda()
noise = torch.randn((3, 120))
res = g(noise.cuda(), self.label.cuda())
assert res.shape == (3, 3, 256, 256)
res = g(None, None, num_batches=3)
assert res.shape == (3, 3, 256, 256)
cfg = deepcopy(self.default_config)
cfg.update(dict(output_scale=512))
g = build_module(cfg).cuda()
res = g(None, None, num_batches=3)
assert res.shape == (3, 3, 512, 512)
cfg = deepcopy(self.default_config)
cfg.update(dict(output_scale=64))
g = build_module(cfg).cuda()
res = g(None, None, num_batches=3)
assert res.shape == (3, 3, 64, 64)
cfg = deepcopy(self.default_config)
cfg.update(dict(output_scale=32))
g = build_module(cfg).cuda()
res = g(None, None, num_batches=3)
assert res.shape == (3, 3, 32, 32)
# test with `concat_noise=False`
cfg = deepcopy(self.default_config)
cfg.update(dict(concat_noise=False))
g = build_module(cfg).cuda()
res = g(None, None, num_batches=3)
assert res.shape == (3, 3, 128, 128)
# test with `with_spectral_norm=False`
cfg = deepcopy(self.default_config)
cfg.update(dict(with_spectral_norm=False))
g = build_module(cfg).cuda()
res = g(None, None, num_batches=3)
assert res.shape == (3, 3, 128, 128)
# test different num_classes
cfg = deepcopy(self.default_config)
cfg.update(
dict(num_classes=0,
with_shared_embedding=False,
concat_noise=False))
g = build_module(cfg).cuda()
res = g(None, None, num_batches=3)
assert res.shape == (3, 3, 128, 128)
# test no shared embedding
cfg = deepcopy(self.default_config)
cfg.update(dict(with_shared_embedding=False, concat_noise=False))
g = build_module(cfg).cuda()
res = g(None, None, num_batches=3)
assert res.shape == (3, 3, 128, 128)
def test_sngan_proj_discriminator_cuda(self):
# test default setting with builder
d = build_module(self.default_config).cuda()
assert isinstance(d, ProjDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
# test different input_scale
config = deepcopy(self.default_config)
config['input_scale'] = 64
d = build_module(config).cuda()
assert isinstance(d, ProjDiscriminator)
x = torch.randn((2, 3, 64, 64)).cuda()
score = d(x, self.label.cuda())
assert score.shape == (2, 1)
# test num_classes == 0 (w/o proj_y)
config = deepcopy(self.default_config)
config['num_classes'] = 0
d = build_module(config).cuda()
assert isinstance(d, ProjDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
# test different base_channels
config = deepcopy(self.default_config)
config['base_channels'] = 128
d = build_module(config).cuda()
assert isinstance(d, ProjDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
# test different channels_cfg --> list
config = deepcopy(self.default_config)
config['channels_cfg'] = [1, 1, 1]
d = build_module(config).cuda()
assert isinstance(d, ProjDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
# test different channels_cfg --> dict
config = deepcopy(self.default_config)
config['channels_cfg'] = {32: [1, 1, 1], 64: [2, 4, 8, 16]}
d = build_module(config).cuda()
assert isinstance(d, ProjDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
# test different downsample_cfg --> list
config = deepcopy(self.default_config)
config['downsample_cfg'] = [True, False, False]
d = build_module(config).cuda()
assert isinstance(d, ProjDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
# test different downsample_cfg --> dict
config = deepcopy(self.default_config)
config['downsample_cfg'] = {
32: [True, False, False],
64: [True, True, True, True]
}
d = build_module(config).cuda()
assert isinstance(d, ProjDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
# test different act_cfg
config = deepcopy(self.default_config)
config['act_cfg'] = dict(type='Sigmoid')
d = build_module(config).cuda()
assert isinstance(d, ProjDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
# test different with_spectral_norm
config = deepcopy(self.default_config)
config['with_spectral_norm'] = False
d = build_module(config).cuda()
assert isinstance(d, ProjDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
# test different init_cfg --> BigGAN
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='biggan')
d = build_module(config).cuda()
assert isinstance(d, ProjDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)
# test different init_cfg --> sngan
config = deepcopy(self.default_config)
config['init_cfg'] = dict(type='sngan-proj')
d = build_module(config).cuda()
assert isinstance(d, ProjDiscriminator)
score = d(self.x.cuda(), self.label.cuda())
assert score.shape == (2, 1)