add tests code

tests/test_datasets/test_unconditional_image_dataset.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import os.path as osp
+
+from mmgen.datasets import UnconditionalImageDataset
+
+
+class TestUnconditionalImageDataset(object):
+
+    @classmethod
+    def setup_class(cls):
+        cls.imgs_root = osp.join(osp.dirname(__file__), '..', 'data/image')
+        cls.default_pipeline = [
+            dict(type='LoadImageFromFile', io_backend='disk', key='real_img')
+        ]
+
+    def test_unconditional_imgs_dataset(self):
+        dataset = UnconditionalImageDataset(
+            self.imgs_root, pipeline=self.default_pipeline)
+        assert len(dataset) == 6
+        img = dataset[2]['real_img']
+        assert img.ndim == 3
+        assert repr(dataset) == (
+            f'dataset_name: {dataset.__class__}, '
+            f'total {6} images in imgs_root: {self.imgs_root}')

tests/test_datasets/test_unpaired_image_dataset.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import os.path as osp
+
+from mmgen.datasets import UnpairedImageDataset
+
+
+class TestUnpairedImageDataset(object):
+
+    @classmethod
+    def setup_class(cls):
+        cls.imgs_root = osp.join(
+            osp.dirname(osp.dirname(__file__)), 'data/unpaired')
+        img_norm_cfg = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        cls.default_pipeline = [
+            dict(
+                type='LoadImageFromFile',
+                io_backend='disk',
+                key='img_a',
+                flag='color'),
+            dict(
+                type='LoadImageFromFile',
+                io_backend='disk',
+                key='img_b',
+                flag='color'),
+            dict(
+                type='Resize',
+                keys=['img_a', 'img_b'],
+                scale=(286, 286),
+                interpolation='bicubic'),
+            dict(
+                type='Crop',
+                keys=['img_a', 'img_b'],
+                crop_size=(256, 256),
+                random_crop=True),
+            dict(type='Flip', keys=['img_a'], direction='horizontal'),
+            dict(type='Flip', keys=['img_b'], direction='horizontal'),
+            dict(type='RescaleToZeroOne', keys=['img_a', 'img_b']),
+            dict(
+                type='Normalize',
+                keys=['img_a', 'img_b'],
+                to_rgb=True,
+                **img_norm_cfg),
+            dict(type='ImageToTensor', keys=['img_a', 'img_b']),
+            dict(
+                type='Collect',
+                keys=['img_a', 'img_b'],
+                meta_keys=['img_a_path', 'img_b_path'])
+        ]
+
+    def test_unpaired_image_dataset(self):
+        dataset = UnpairedImageDataset(
+            self.imgs_root,
+            pipeline=self.default_pipeline,
+            domain_a='a',
+            domain_b='b')
+        assert len(dataset) == 2
+        img = dataset[0]['img_a']
+        assert img.ndim == 3
+        img = dataset[0]['img_b']
+        assert img.ndim == 3

tests/test_losses/test_ddpm_loss.py

+# Copyright (c) OpenMMLab. All rights reserved.
+from copy import deepcopy
+
+import numpy as np
+import pytest
+import torch
+
+from mmgen.models.builder import build_module
+from mmgen.models.losses.pixelwise_loss import (
+    DiscretizedGaussianLogLikelihoodLoss, GaussianKLDLoss, MSELoss)
+
+
+class TestDDPMVLBLoss:
+
+    @classmethod
+    def setup_class(cls):
+        cls.gaussian_kld_data_info = dict(
+            mean_pred='mean_pred',
+            mean_target='mean_posterior',
+            logvar_pred='logvar_pred',
+            logvar_target='logvar_posterior')
+        cls.disc_log_likelihood_data_info = dict(
+            x='real_imgs', mean='mean_pred', logvar='logvar_pred')
+        cls.config = dict(
+            type='DDPMVLBLoss',
+            rescale_mode='constant',
+            rescale_cfg=dict(scale=4),
+            data_info=cls.gaussian_kld_data_info,
+            data_info_t_0=cls.disc_log_likelihood_data_info,
+            log_cfgs=[
+                dict(
+                    type='quartile', prefix_name='loss_vlb',
+                    total_timesteps=4),
+                dict(type='name')
+            ])
+        cls.t = torch.LongTensor([0, 1, 2, 3])
+        cls.tar_shape = [4, 2, 4, 4]
+
+        cls.mean_pred = torch.randn(cls.tar_shape)
+        cls.logvar_pred = torch.randn(cls.tar_shape)
+        cls.mean_posterior = torch.randn(cls.tar_shape)
+        cls.logvar_posterior = torch.randn(cls.tar_shape)
+        cls.real_imgs = torch.randn(cls.tar_shape)
+        cls.label = [0, 18, 1, 5]
+
+        cls.output_dict = dict(
+            mean_pred=cls.mean_pred,
+            logvar_pred=cls.logvar_pred,
+            mean_posterior=cls.mean_posterior,
+            logvar_posterior=cls.logvar_posterior,
+            real_imgs=cls.real_imgs,
+            label=cls.label,
+            meta_info=None,
+            timesteps=cls.t)
+
+        # calculate loss manually
+        cls.loss_gaussian_kld = GaussianKLDLoss(
+            data_info=cls.gaussian_kld_data_info,
+            reduction='flatmean',
+            base='2')(
+                cls.output_dict)
+        cls.loss_disc_likelihood = DiscretizedGaussianLogLikelihoodLoss(
+            data_info=cls.disc_log_likelihood_data_info,
+            reduction='flatmean',
+            base='2')(
+                cls.output_dict)
+
+        cls.loss_manually = (-cls.loss_disc_likelihood[0] +
+                             cls.loss_gaussian_kld[1:].sum()) / 4
+
+        # TODO: unit test for sampler would be add later
+        cls.weight = torch.rand(4, )
+
+    def test_vlb_loss(self):
+        # test forward
+        config = deepcopy(self.config)
+        loss_fn = build_module(config)
+        loss = loss_fn(self.output_dict)
+        np.allclose(loss, self.loss_manually * 4)
+
+        # test log_cfgs --> dict input
+        config = deepcopy(self.config)
+        config['log_cfgs'] = dict(type='name')
+        loss_fn = build_module(config)
+        assert isinstance(loss_fn.log_fn_list, list)
+
+        # test log_cfgs --> no log_cfgs
+        config = deepcopy(self.config)
+        config['log_cfgs'] = None
+        loss_fn = build_module(config)
+        loss = loss_fn(self.output_dict)
+        assert not loss_fn.log_vars
+
+        # test rescale_cfg --> rescale is None
+        config = deepcopy(self.config)
+        config['rescale_mode'] = None
+        loss_fn = build_module(config)
+        loss = loss_fn(self.output_dict)
+        np.allclose(loss, self.loss_manually)
+
+        # TODO: test rescale_cfg --> test sampler
+
+        # test rescale_cfg --> test weight
+        config = deepcopy(self.config)
+        config['rescale_mode'] = 'timestep_weight'
+        weight = self.weight.clone()
+        loss_fn = build_module(config, default_args=dict(weight=weight))
+        loss = loss_fn(self.output_dict)
+        loss_weighted_manually = (
+            -(self.loss_disc_likelihood * weight)[0] +
+            (self.loss_gaussian_kld * weight)[1:].sum()) / 4
+        np.allclose(loss, loss_weighted_manually)
+
+        # test rescale_cfg --> change weight
+        weight[0] += 1
+        loss = loss_fn(self.output_dict)
+        loss_weighted_manually = (
+            -(self.loss_disc_likelihood * weight)[0] +
+            (self.loss_gaussian_kld * weight)[1:].sum()) / 4
+        np.allclose(loss, loss_weighted_manually)
+
+        # test t = 0
+        config = deepcopy(self.config)
+        output_dict = deepcopy(self.output_dict)
+        output_dict['timesteps'][0] = 1
+        loss_fn = build_module(config)
+        loss = loss_fn(output_dict)
+        assert loss_fn.log_vars['loss_vlb_quartile_0'] == 0
+        assert loss_fn.log_vars['loss_DiscGaussianLogLikelihood'] == 0
+
+
+class TestDDPMMSELoss:
+
+    @classmethod
+    def setup_class(cls):
+        cls.data_info = dict(pred='eps_t_pred', target='noise')
+        cls.config = dict(
+            type='DDPMMSELoss',
+            data_info=cls.data_info,
+            log_cfgs=dict(
+                type='quartile', prefix_name='loss_mse', total_timesteps=4))
+
+        cls.t = torch.LongTensor([0, 1, 2, 3])
+        cls.tar_shape = [4, 2, 4, 4]
+
+        cls.eps_t_pred = torch.randn(cls.tar_shape)
+        cls.noise = torch.randn(cls.tar_shape)
+
+        cls.output_dict = dict(
+            eps_t_pred=cls.eps_t_pred,
+            noise=cls.noise,
+            meta_info=None,
+            timesteps=cls.t)
+
+        cls.weight = torch.rand(4, )
+
+        # calculate loss manually
+        cls.loss_manually = 0
+        for idx in range(cls.tar_shape[0]):
+            t = cls.t[idx]
+            weight = cls.weight[t]
+            output_dict_ = dict(
+                eps_t_pred=cls.eps_t_pred[t], noise=cls.noise[t])
+            cls.loss_manually += MSELoss(
+                data_info=cls.data_info)(output_dict_) * weight
+        cls.loss_manually /= 4
+
+    def test_mse_loss(self):
+        # test forward
+        config = deepcopy(self.config)
+        config['rescale_mode'] = 'timestep_weight'
+        loss_fn = build_module(config, default_args=dict(weight=self.weight))
+        loss = loss_fn(self.output_dict)
+        np.allclose(loss, self.loss_manually)
+
+        # test reduction raise error
+        config = deepcopy(self.config)
+        config['reduction'] = 'reduction'
+        with pytest.raises(ValueError):
+            loss_fn = build_module(config)
+
+        # test return loss name
+        config = deepcopy(self.config)
+        config['loss_name'] = 'loss_name'
+        loss_fn = build_module(config)
+        assert loss_fn.loss_name() == 'loss_name'

tests/test_losses/test_disc_auxilary_loss.py

+# Copyright (c) OpenMMLab. All rights reserved.
+from functools import partial
+
+import pytest
+import torch
+
+from mmgen.models.architectures import DCGANDiscriminator
+from mmgen.models.architectures.pggan.generator_discriminator import \
+    PGGANDiscriminator
+from mmgen.models.losses import (DiscShiftLoss, GradientPenaltyLoss,
+                                 disc_shift_loss, gradient_penalty_loss)
+from mmgen.models.losses.disc_auxiliary_loss import (R1GradientPenalty,
+                                                     r1_gradient_penalty_loss)
+
+
+class TestDiscShiftLoss(object):
+
+    @classmethod
+    def setup_class(cls):
+        cls.input_tensor = torch.randn((2, 10))
+        cls.default_cfg = dict(
+            loss_weight=0.1, data_info=dict(pred='disc_pred'))
+        cls.default_input_dict = dict(disc_pred=cls.input_tensor)
+
+    def test_disc_shift_loss(self):
+        loss = disc_shift_loss(self.input_tensor)
+        assert loss.ndim == 0
+        assert loss >= 0
+
+        loss = disc_shift_loss(self.input_tensor, weight=-0.1)
+        assert loss.ndim == 0
+        assert loss <= 0
+
+        loss = disc_shift_loss(self.input_tensor, reduction='none')
+        assert loss.ndim == 2
+        assert (loss >= 0).all()
+
+        loss_sum = disc_shift_loss(self.input_tensor, reduction='sum')
+        loss_avg = disc_shift_loss(self.input_tensor, avg_factor=1000)
+        assert loss_avg.ndim == 0 and loss_sum.ndim == 0
+        assert loss_sum > loss_avg
+
+        with pytest.raises(ValueError):
+            _ = disc_shift_loss(
+                self.input_tensor, reduction='sum', avg_factor=100)
+
+    def test_module_wrapper(self):
+        # test with default config
+        loss_module = DiscShiftLoss(**self.default_cfg)
+        loss = loss_module(self.default_input_dict)
+        assert loss.ndim == 0
+
+        with pytest.raises(NotImplementedError):
+            _ = loss_module(self.default_input_dict, 1)
+
+        with pytest.raises(AssertionError):
+            _ = loss_module(1, outputs_dict=self.default_input_dict)
+        input_ = dict(outputs_dict=self.default_input_dict)
+        loss = loss_module(**input_)
+        assert loss.ndim == 0
+
+        with pytest.raises(AssertionError):
+            _ = loss_module(self.input_tensor)
+
+        # test without data_info
+        loss_module = DiscShiftLoss(data_info=None)
+        loss = loss_module(self.input_tensor)
+        assert loss.ndim == 0
+
+
+class TestGradientPenalty:
+
+    @classmethod
+    def setup_class(cls):
+        cls.input_img = torch.randn((2, 3, 8, 8))
+        cls.disc = DCGANDiscriminator(
+            input_scale=8, output_scale=4, out_channels=5)
+        cls.pggan_disc = PGGANDiscriminator(
+            in_scale=8, base_channels=32, max_channels=32)
+        cls.data_info = dict(
+            discriminator='disc', real_data='real_imgs', fake_data='fake_imgs')
+
+    def test_gp_loss(self):
+        loss = gradient_penalty_loss(self.disc, self.input_img, self.input_img)
+        assert loss > 0
+
+        loss = gradient_penalty_loss(
+            self.disc, self.input_img, self.input_img, norm_mode='HWC')
+        assert loss > 0
+
+        with pytest.raises(NotImplementedError):
+            _ = gradient_penalty_loss(
+                self.disc, self.input_img, self.input_img, norm_mode='xxx')
+
+        loss = gradient_penalty_loss(
+            self.disc,
+            self.input_img,
+            self.input_img,
+            norm_mode='HWC',
+            weight=10)
+        assert loss > 0
+
+        loss = gradient_penalty_loss(
+            self.disc,
+            self.input_img,
+            self.input_img,
+            norm_mode='HWC',
+            mask=torch.ones_like(self.input_img),
+            weight=10)
+        assert loss > 0
+
+        data_dict = dict(
+            real_imgs=self.input_img,
+            fake_imgs=self.input_img,
+            disc=partial(self.pggan_disc, transition_weight=0.5, curr_scale=8))
+        gp_loss = GradientPenaltyLoss(
+            loss_weight=10, norm_mode='pixel', data_info=self.data_info)
+
+        loss = gp_loss(data_dict)
+        assert loss > 0
+        loss = gp_loss(outputs_dict=data_dict)
+        assert loss > 0
+        with pytest.raises(NotImplementedError):
+            _ = gp_loss(asdf=1.)
+
+        with pytest.raises(AssertionError):
+            _ = gp_loss(1.)
+
+        with pytest.raises(AssertionError):
+            _ = gp_loss(1., 2, outputs_dict=data_dict)
+
+
+class TestR1GradientPenalty:
+
+    @classmethod
+    def setup_class(cls):
+        cls.data_info = dict(discriminator='disc', real_data='real_imgs')
+        cls.disc = DCGANDiscriminator(
+            input_scale=8, output_scale=4, out_channels=5)
+        cls.pggan_disc = PGGANDiscriminator(
+            in_scale=8, base_channels=32, max_channels=32)
+        cls.input_img = torch.randn((2, 3, 8, 8))
+
+    def test_r1_regularizer(self):
+        loss = r1_gradient_penalty_loss(self.disc, self.input_img)
+        assert loss > 0
+
+        loss = r1_gradient_penalty_loss(
+            self.disc, self.input_img, norm_mode='HWC')
+        assert loss > 0
+
+        with pytest.raises(NotImplementedError):
+            _ = r1_gradient_penalty_loss(
+                self.disc, self.input_img, norm_mode='xxx')
+
+        loss = r1_gradient_penalty_loss(
+            self.disc, self.input_img, norm_mode='HWC', weight=10)
+        assert loss > 0
+
+        loss = r1_gradient_penalty_loss(
+            self.disc,
+            self.input_img,
+            norm_mode='HWC',
+            mask=torch.ones_like(self.input_img),
+            weight=10)
+        assert loss > 0
+
+        data_dict = dict(
+            real_imgs=self.input_img,
+            disc=partial(self.pggan_disc, transition_weight=0.5, curr_scale=8))
+        gp_loss = R1GradientPenalty(
+            loss_weight=10, norm_mode='pixel', data_info=self.data_info)
+
+        loss = gp_loss(data_dict)
+        assert loss > 0
+        loss = gp_loss(outputs_dict=data_dict)
+        assert loss > 0
+        with pytest.raises(NotImplementedError):
+            _ = gp_loss(asdf=1.)
+
+        with pytest.raises(AssertionError):
+            _ = gp_loss(1.)
+
+        with pytest.raises(AssertionError):
+            _ = gp_loss(1., 2, outputs_dict=data_dict)

tests/test_losses/test_gan_loss.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import numpy.testing as npt
+import pytest
+import torch
+
+from mmgen.models.losses.gan_loss import GANLoss
+
+
+def test_gan_losses():
+    """Test gan losses."""
+    with pytest.raises(NotImplementedError):
+        GANLoss(
+            'xixihaha',
+            loss_weight=1.0,
+            real_label_val=1.0,
+            fake_label_val=0.0)
+
+    input_1 = torch.ones(1, 1)
+    input_2 = torch.ones(1, 3, 6, 6) * 2
+
+    # vanilla
+    gan_loss = GANLoss(
+        'vanilla', loss_weight=2.0, real_label_val=1.0, fake_label_val=0.0)
+    loss = gan_loss(input_1, True, is_disc=False)
+    npt.assert_almost_equal(loss.item(), 0.6265233)
+    loss = gan_loss(input_1, False, is_disc=False)
+    npt.assert_almost_equal(loss.item(), 2.6265232)
+    loss = gan_loss(input_1, True, is_disc=True)
+    npt.assert_almost_equal(loss.item(), 0.3132616)
+    loss = gan_loss(input_1, False, is_disc=True)
+    npt.assert_almost_equal(loss.item(), 1.3132616)
+
+    # lsgan
+    gan_loss = GANLoss(
+        'lsgan', loss_weight=2.0, real_label_val=1.0, fake_label_val=0.0)
+    loss = gan_loss(input_2, True, is_disc=False)
+    npt.assert_almost_equal(loss.item(), 2.0)
+    loss = gan_loss(input_2, False, is_disc=False)
+    npt.assert_almost_equal(loss.item(), 8.0)
+    loss = gan_loss(input_2, True, is_disc=True)
+    npt.assert_almost_equal(loss.item(), 1.0)
+    loss = gan_loss(input_2, False, is_disc=True)
+    npt.assert_almost_equal(loss.item(), 4.0)
+
+    # wgan
+    gan_loss = GANLoss(
+        'wgan', loss_weight=2.0, real_label_val=1.0, fake_label_val=0.0)
+    loss = gan_loss(input_2, True, is_disc=False)
+    npt.assert_almost_equal(loss.item(), -4.0)
+    loss = gan_loss(input_2, False, is_disc=False)
+    npt.assert_almost_equal(loss.item(), 4)
+    loss = gan_loss(input_2, True, is_disc=True)
+    npt.assert_almost_equal(loss.item(), -2.0)
+    loss = gan_loss(input_2, False, is_disc=True)
+    npt.assert_almost_equal(loss.item(), 2.0)
+
+    # wgan
+    gan_loss = GANLoss(
+        'wgan-logistic-ns',
+        loss_weight=2.0,
+        real_label_val=1.0,
+        fake_label_val=0.0)
+    loss = gan_loss(input_2, True, is_disc=False)
+    assert loss.item() > 0
+    loss = gan_loss(input_2, False, is_disc=False)
+    assert loss.item() > 0
+
+    # hinge
+    gan_loss = GANLoss(
+        'hinge', loss_weight=2.0, real_label_val=1.0, fake_label_val=0.0)
+    loss = gan_loss(input_2, True, is_disc=False)
+    npt.assert_almost_equal(loss.item(), -4.0)
+    loss = gan_loss(input_2, False, is_disc=False)
+    npt.assert_almost_equal(loss.item(), -4.0)
+    loss = gan_loss(input_2, True, is_disc=True)
+    npt.assert_almost_equal(loss.item(), 0.0)
+    loss = gan_loss(input_2, False, is_disc=True)
+    npt.assert_almost_equal(loss.item(), 3.0)

tests/test_losses/test_gen_auxiliary_loss.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import pytest
+import torch
+
+from mmgen.models.architectures.pix2pix import UnetGenerator
+from mmgen.models.architectures.stylegan import StyleGANv2Generator
+from mmgen.models.losses import GeneratorPathRegularizer, PerceptualLoss
+from mmgen.models.losses.pixelwise_loss import l1_loss, mse_loss
+
+
+class TestPathRegularizer:
+
+    @classmethod
+    def setup_class(cls):
+        cls.data_info = dict(generator='generator', num_batches='num_batches')
+        cls.gen = StyleGANv2Generator(32, 10, num_mlps=2)
+
+    def test_path_regularizer_cpu(self):
+        gen = self.gen
+
+        output_dict = dict(generator=gen, num_batches=2)
+        pl = GeneratorPathRegularizer(data_info=self.data_info)
+        pl_loss = pl(output_dict)
+        assert pl_loss > 0
+
+        output_dict = dict(generator=gen, num_batches=2, iteration=3)
+        pl = GeneratorPathRegularizer(data_info=self.data_info, interval=2)
+        pl_loss = pl(outputs_dict=output_dict)
+        assert pl_loss is None
+
+        with pytest.raises(NotImplementedError):
+            _ = pl(asdf=1.)
+
+        with pytest.raises(AssertionError):
+            _ = pl(1.)
+
+        with pytest.raises(AssertionError):
+            _ = pl(1., 2, outputs_dict=output_dict)
+
+    @pytest.mark.skipif(
+        not torch.cuda.is_available()
+        or not hasattr(torch.backends.cudnn, 'allow_tf32'),
+        reason='requires cuda')
+    def test_path_regularizer_cuda(self):
+        gen = self.gen.cuda()
+
+        output_dict = dict(generator=gen, num_batches=2)
+        pl = GeneratorPathRegularizer(data_info=self.data_info).cuda()
+        pl_loss = pl(output_dict)
+        assert pl_loss > 0
+
+        output_dict = dict(generator=gen, num_batches=2, iteration=3)
+        pl = GeneratorPathRegularizer(
+            data_info=self.data_info, interval=2).cuda()
+        pl_loss = pl(outputs_dict=output_dict)
+        assert pl_loss is None
+
+        with pytest.raises(NotImplementedError):
+            _ = pl(asdf=1.)
+
+        with pytest.raises(AssertionError):
+            _ = pl(1.)
+
+        with pytest.raises(AssertionError):
+            _ = pl(1., 2, outputs_dict=output_dict)
+
+
+class TestPerceptualLoss:
+
+    @classmethod
+    def setup_class(cls):
+        cls.data_info = dict(pred='fake_imgs', target='real_imgs')
+        cls.gen = UnetGenerator(3, 3)
+
+    def test_perceptual_loss_cpu(self):
+        unknown_h, unknown_w = (32, 32)
+        weight = torch.zeros(1, 1, 64, 64)
+        weight[0, 0, :unknown_h, :unknown_w] = 1
+        pred = weight.clone()
+        target = weight.clone() * 2
+
+        perceptual_loss = PerceptualLoss(data_info=self.data_info)
+        loss_perceptual = perceptual_loss(
+            outputs_dict=dict(fake_imgs=pred, real_imgs=target))
+        assert loss_perceptual.shape == ()
+        assert id(perceptual_loss.criterion) == id(l1_loss)
+
+    def test_only_perceptual_loss(self):
+        unknown_h, unknown_w = (32, 32)
+        weight = torch.zeros(1, 1, 64, 64)
+        weight[0, 0, :unknown_h, :unknown_w] = 1
+        pred = weight.clone()
+        target = weight.clone() * 2
+
+        perceptual_loss = PerceptualLoss(
+            data_info=self.data_info, style_weight=0)
+        loss_percep = perceptual_loss(dict(fake_imgs=pred, real_imgs=target))
+        assert loss_percep.shape == ()
+        assert perceptual_loss.style_weight == 0
+
+    def test_only_style_loss(self):
+        unknown_h, unknown_w = (32, 32)
+        weight = torch.zeros(1, 1, 64, 64)
+        weight[0, 0, :unknown_h, :unknown_w] = 1
+        pred = weight.clone()
+        target = weight.clone() * 2
+
+        perceptual_loss = PerceptualLoss(
+            data_info=self.data_info, perceptual_weight=0)
+        loss_style = perceptual_loss(dict(fake_imgs=pred, real_imgs=target))
+        assert loss_style.shape == ()
+        assert perceptual_loss.perceptual_weight == 0
+
+    def test_with_different_layer_weights(self):
+        unknown_h, unknown_w = (32, 32)
+        weight = torch.zeros(1, 1, 64, 64)
+        weight[0, 0, :unknown_h, :unknown_w] = 1
+        pred = weight.clone()
+        target = weight.clone() * 2
+
+        layer_weights = {'1': 1., '2': 2., '3': 3.}
+        perceptual_loss = PerceptualLoss(
+            data_info=self.data_info, layer_weights=layer_weights)
+        loss_perceptual = perceptual_loss(
+            dict(fake_imgs=pred, real_imgs=target))
+        assert loss_perceptual.shape == ()
+        assert perceptual_loss.layer_weights == layer_weights and \
+            perceptual_loss.layer_weights_style == layer_weights
+
+    def test_with_different_perceptual_and_style_layers(self):
+        unknown_h, unknown_w = (32, 32)
+        weight = torch.zeros(1, 1, 64, 64)
+        weight[0, 0, :unknown_h, :unknown_w] = 1
+        pred = weight.clone()
+        target = weight.clone() * 2
+
+        layer_weights = {'1': 1., '2': 2., '3': 3.}
+        layer_weights_style = {'4': 4., '5': 5., '6': 6.}
+        perceptual_loss = PerceptualLoss(
+            data_info=self.data_info,
+            layer_weights=layer_weights,
+            layer_weights_style=layer_weights_style)
+        loss_perceptual = perceptual_loss(
+            dict(fake_imgs=pred, real_imgs=target))
+        assert loss_perceptual.shape == ()
+        assert perceptual_loss.layer_weights == layer_weights and \
+            perceptual_loss.layer_weights_style == layer_weights_style
+
+    def test_MSE_critierion(self):
+        unknown_h, unknown_w = (32, 32)
+        weight = torch.zeros(1, 1, 64, 64)
+        weight[0, 0, :unknown_h, :unknown_w] = 1
+        pred = weight.clone()
+        target = weight.clone() * 2
+
+        perceptual_loss = PerceptualLoss(
+            data_info=self.data_info, criterion='mse')
+        loss_perceptual = perceptual_loss(
+            outputs_dict=dict(fake_imgs=pred, real_imgs=target))
+        assert loss_perceptual.shape == ()
+        assert id(perceptual_loss.criterion) == id(mse_loss)
+
+    def test_VGG_16(self):
+        unknown_h, unknown_w = (32, 32)
+        weight = torch.zeros(1, 1, 64, 64)
+        weight[0, 0, :unknown_h, :unknown_w] = 1
+        pred = weight.clone()
+        target = weight.clone() * 2
+
+        perceptual_loss = PerceptualLoss(
+            data_info=self.data_info,
+            vgg_type='vgg16',
+            pretrained='torchvision://vgg16')
+        loss_perceptual = perceptual_loss(
+            outputs_dict=dict(fake_imgs=pred, real_imgs=target))
+        assert loss_perceptual.shape == ()
+        # TODO need to check whether vgg16 is loaded
+        # assert perceptual_loss.vgg
+
+    def test_split_style_loss(self):
+        unknown_h, unknown_w = (32, 32)
+        weight = torch.zeros(1, 1, 64, 64)
+        weight[0, 0, :unknown_h, :unknown_w] = 1
+        pred = weight.clone()
+        target = weight.clone() * 2
+
+        perceptual_loss = PerceptualLoss(
+            data_info=self.data_info, split_style_loss=True)
+        loss_percep, loss_style = perceptual_loss(
+            outputs_dict=dict(fake_imgs=pred, real_imgs=target))
+        assert loss_percep.shape == () and loss_style.shape == ()
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda')
+    def test_perceptual_loss_cuda(self):
+        pred = torch.rand([2, 3, 256, 256]).cuda()
+        target = torch.rand_like(pred).cuda()
+        perceptual_loss = PerceptualLoss(data_info=self.data_info).cuda()
+        loss_perceptual = perceptual_loss(
+            outputs_dict=dict(fake_imgs=pred, real_imgs=target))
+        assert loss_perceptual.shape == ()

tests/test_losses/test_pixelwise_loss.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import pytest
+import torch
+from torch.distributions.normal import Normal
+
+from mmgen.models.architectures.pix2pix import UnetGenerator
+from mmgen.models.losses import L1Loss, MSELoss
+from mmgen.models.losses.pixelwise_loss import (
+    DiscretizedGaussianLogLikelihoodLoss, GaussianKLDLoss, approx_gaussian_cdf)
+
+
+class TestPixelwiseLosses:
+
+    @classmethod
+    def setup_class(cls):
+        cls.data_info = dict(pred='fake_imgs', target='real_imgs')
+        cls.gen = UnetGenerator(3, 3)
+
+    def test_pixelwise_losses(self):
+        with pytest.raises(ValueError):
+            # only 'none', 'mean' and 'sum' are supported
+            L1Loss(reduction='InvalidValue')
+
+        unknown_h, unknown_w = (32, 32)
+        weight = torch.zeros(1, 1, 64, 64)
+        weight[0, 0, :unknown_h, :unknown_w] = 1
+        pred = weight.clone()
+        target = weight.clone() * 2
+
+        # test l1 loss
+        l1_loss = L1Loss(
+            loss_weight=1.0, reduction='mean', data_info=self.data_info)
+        loss = l1_loss(outputs_dict=dict(fake_imgs=pred, real_imgs=target))
+        assert loss.shape == ()
+        assert loss.item() == 0.25
+
+        l1_loss = L1Loss(
+            loss_weight=1.0, reduction='mean', data_info=self.data_info)
+        loss = l1_loss(dict(fake_imgs=pred, real_imgs=target))
+        assert loss.shape == ()
+        assert loss.item() == 0.25
+
+        l1_loss = L1Loss(loss_weight=0.5, reduction='none')
+        loss = l1_loss(pred, target)
+        assert loss.shape == (1, 1, 64, 64)
+        assert (loss == torch.ones(1, 1, 64, 64) * weight * 0.5).all()
+
+        l1_loss = L1Loss(loss_weight=0.5, reduction='sum')
+        loss = l1_loss(pred, target)
+        assert loss.shape == ()
+        assert loss.item() == 512
+
+        # test MSE loss
+        mse_loss = MSELoss(loss_weight=1.0, data_info=self.data_info)
+        loss = mse_loss(outputs_dict=dict(fake_imgs=pred, real_imgs=target))
+        assert loss.shape == ()
+        assert loss.item() == 0.25
+
+        mse_loss = MSELoss(loss_weight=1.0, data_info=self.data_info)
+        loss = mse_loss(dict(fake_imgs=pred, real_imgs=target))
+        assert loss.shape == ()
+        assert loss.item() == 0.25
+
+        mse_loss = MSELoss(loss_weight=0.5)
+        loss = mse_loss(pred, target)
+        assert loss.shape == ()
+        assert loss.item() == 0.1250
+
+
+class TestGaussianKLDLoss:
+
+    @classmethod
+    def setup_class(cls):
+        cls.data_info = dict(
+            mean_pred='mean_pred',
+            mean_target='mean_target',
+            logvar_pred='logvar_pred',
+            logvar_target='logvar_target')
+        cls.tar_shape = [2, 2, 4, 4]
+
+        cls.mean_pred = torch.zeros(cls.tar_shape)
+        cls.mean_target = torch.ones(cls.tar_shape)
+        cls.logvar_pred = torch.zeros(cls.tar_shape)
+        cls.logvar_target = torch.ones(cls.tar_shape)
+        cls.output_dict = dict(
+            mean_pred=cls.mean_pred,
+            mean_target=cls.mean_target,
+            logvar_pred=cls.logvar_pred,
+            logvar_target=cls.logvar_target)
+        cls.gt_loss = ((torch.exp(torch.ones(1)) - 1) / 2).item()
+
+    def test_gaussian_kld_loss(self):
+
+        # test reduction --> mean
+        gaussian_kld_loss = GaussianKLDLoss(
+            data_info=self.data_info, reduction='mean')
+        loss = gaussian_kld_loss(self.output_dict)
+        assert (loss == self.gt_loss).all()
+
+        # test reduction --> batchmean
+        gaussian_kld_loss = GaussianKLDLoss(
+            data_info=self.data_info, reduction='batchmean')
+        loss = gaussian_kld_loss(self.output_dict)
+        num_elements = self.tar_shape[1] * self.tar_shape[2] * \
+            self.tar_shape[3]
+        assert (loss == (self.gt_loss * num_elements)).all()
+
+        # test weight --> int
+        gaussian_kld_loss = GaussianKLDLoss(
+            loss_weight=2, data_info=self.data_info, reduction='mean')
+        loss = gaussian_kld_loss(self.output_dict)
+        assert (loss == self.gt_loss * 2).all()
+
+        # test weight --> tensor
+        weight = torch.randn(*self.tar_shape)
+        gaussian_kld_loss = GaussianKLDLoss(
+            loss_weight=weight, data_info=self.data_info, reduction='mean')
+        loss = gaussian_kld_loss(self.output_dict)
+        assert torch.allclose(loss, weight.mean() * self.gt_loss, atol=1e-6)
+
+        # test weight --> tensor & batchmean
+        weight = torch.randn(*self.tar_shape)
+        gaussian_kld_loss = GaussianKLDLoss(
+            loss_weight=weight,
+            data_info=self.data_info,
+            reduction='batchmean')
+        loss = gaussian_kld_loss(self.output_dict)
+        assert torch.allclose(
+            loss, weight.sum([1, 2, 3]).mean() * self.gt_loss, atol=1e-6)
+
+
+def test_approx_gaussian_cdf():
+    pos = torch.rand(2, 2)
+    gaussian_dist = Normal(0, 1)
+    assert torch.allclose(
+        approx_gaussian_cdf(pos), gaussian_dist.cdf(pos), atol=1e-3)
+
+
+class TestDistLoss():
+
+    @classmethod
+    def setup_class(cls):
+        cls.data_info = dict(
+            mean='mean_pred', logvar='logvar_pred', x='real_imgs')
+        cls.tar_shape = [2, 2, 4, 4]
+
+        cls.mean_pred = torch.zeros(cls.tar_shape)
+        cls.logvar_pred = torch.zeros(cls.tar_shape)
+        cls.real_imgs = torch.zeros(cls.tar_shape)
+
+        cls.output_dict = dict(
+            mean_pred=cls.mean_pred,
+            logvar_pred=cls.logvar_pred,
+            real_imgs=cls.real_imgs)
+        norm_dist = Normal(0, 1)
+        cls.gt_loss = torch.log(
+            norm_dist.cdf(torch.FloatTensor([1 / 255])) -
+            norm_dist.cdf(torch.FloatTensor([-1 / 255])))
+
+    def test_disc_gaussian_log_likelihood_loss(self):
+        # test reduction --> mean
+        disc_gaussian_loss = DiscretizedGaussianLogLikelihoodLoss(
+            data_info=self.data_info, reduction='mean')
+        loss = disc_gaussian_loss(self.output_dict)
+        assert (loss == self.gt_loss).all()
+
+        # test reduction --> batchmean
+        disc_gaussian_loss = DiscretizedGaussianLogLikelihoodLoss(
+            data_info=self.data_info, reduction='batchmean')
+        loss = disc_gaussian_loss(self.output_dict)
+        num_elements = self.tar_shape[1] * self.tar_shape[2] * \
+            self.tar_shape[3]
+        assert (loss == (self.gt_loss * num_elements)).all()
+
+        # test weight --> int
+        disc_gaussian_loss = DiscretizedGaussianLogLikelihoodLoss(
+            loss_weight=2, data_info=self.data_info, reduction='mean')
+        loss = disc_gaussian_loss(self.output_dict)
+        assert (loss == self.gt_loss * 2).all()
+
+        # # test weight --> tensor
+        weight = torch.randn(*self.tar_shape)
+        disc_gaussian_loss = DiscretizedGaussianLogLikelihoodLoss(
+            loss_weight=weight, data_info=self.data_info, reduction='mean')
+        loss = disc_gaussian_loss(self.output_dict)
+        assert torch.allclose(loss, weight.mean() * self.gt_loss, atol=1e-6)
+
+        # test weight --> tensor & batchmean
+        weight = torch.randn(*self.tar_shape)
+        disc_gaussian_loss = DiscretizedGaussianLogLikelihoodLoss(
+            loss_weight=weight,
+            data_info=self.data_info,
+            reduction='batchmean')
+        loss = disc_gaussian_loss(self.output_dict)
+        assert torch.allclose(
+            loss, weight.sum([1, 2, 3]).mean() * self.gt_loss, atol=1e-6)

tests/test_models/test_base_ddpm.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+
+from mmgen.models.diffusions import UniformTimeStepSampler
+
+
+def test_uniform_sampler():
+    sampler = UniformTimeStepSampler(10)
+    timesteps = sampler(2)
+    assert timesteps.shape == torch.Size([
+        2,
+    ])
+    assert timesteps.max() < 10 and timesteps.min() >= 0
+
+    timesteps = sampler.__call__(2)
+    assert timesteps.shape == torch.Size([
+        2,
+    ])
+    assert timesteps.max() < 10 and timesteps.min() >= 0

tests/test_models/test_base_gan.py

+# Copyright (c) OpenMMLab. All rights reserved.

tests/test_models/test_basic_conditional_gan.py

+# Copyright (c) OpenMMLab. All rights reserved.
+from copy import deepcopy
+
+import pytest
+import torch
+import torch.nn as nn
+
+from mmgen.models import BasicConditionalGAN, build_model
+
+
+class TestBasicConditionalGAN(object):
+
+    @classmethod
+    def setup_class(cls):
+        cls.default_config = dict(
+            type='BasicConditionalGAN',
+            generator=dict(
+                type='SNGANGenerator',
+                output_scale=32,
+                base_channels=256,
+                num_classes=10),
+            discriminator=dict(
+                type='ProjDiscriminator',
+                input_scale=32,
+                base_channels=128,
+                num_classes=10),
+            gan_loss=dict(type='GANLoss', gan_type='hinge'),
+            disc_auxiliary_loss=None,
+            gen_auxiliary_loss=None,
+            train_cfg=None,
+            test_cfg=None)
+
+        cls.generator_cfg = dict(
+            type='SAGANGenerator',
+            output_scale=32,
+            num_classes=10,
+            base_channels=256,
+            attention_after_nth_block=2,
+            with_spectral_norm=True)
+        cls.disc_cfg = dict(
+            type='SAGANDiscriminator',
+            input_scale=32,
+            num_classes=10,
+            base_channels=128,
+            attention_after_nth_block=1)
+        cls.gan_loss = dict(type='GANLoss', gan_type='hinge')
+        cls.disc_auxiliary_loss = [
+            dict(
+                type='DiscShiftLoss',
+                loss_weight=0.5,
+                data_info=dict(pred='disc_pred_fake')),
+            dict(
+                type='DiscShiftLoss',
+                loss_weight=0.5,
+                data_info=dict(pred='disc_pred_real'))
+        ]
+
+    def test_default_dcgan_model_cpu(self):
+        sngan = build_model(self.default_config)
+        assert isinstance(sngan, BasicConditionalGAN)
+        assert not sngan.with_disc_auxiliary_loss
+        assert sngan.with_disc
+
+        # test forward train
+        with pytest.raises(NotImplementedError):
+            _ = sngan(None, return_loss=True)
+        # test forward test
+        imgs = sngan(None, return_loss=False, mode='sampling', num_batches=2)
+        assert imgs.shape == (2, 3, 32, 32)
+
+        # test train step
+        data = torch.randn((2, 3, 32, 32))
+        label = torch.randint(0, 10, (2, ))
+        data_input = dict(img=data, gt_label=label)
+        optimizer_g = torch.optim.SGD(sngan.generator.parameters(), lr=0.01)
+        optimizer_d = torch.optim.SGD(
+            sngan.discriminator.parameters(), lr=0.01)
+        optim_dict = dict(generator=optimizer_g, discriminator=optimizer_d)
+
+        model_outputs = sngan.train_step(data_input, optim_dict)
+        assert 'results' in model_outputs
+        assert 'log_vars' in model_outputs
+        assert model_outputs['num_samples'] == 2
+
+        # more tests for different configs with heavy computation
+        # test disc_steps
+        config_ = deepcopy(self.default_config)
+        config_['train_cfg'] = dict(disc_steps=2)
+        sngan = build_model(config_)
+        model_outputs = sngan.train_step(data_input, optim_dict)
+        assert 'loss_disc_fake' in model_outputs['log_vars']
+        assert 'loss_disc_fake_g' not in model_outputs['log_vars']
+        assert sngan.disc_steps == 2
+
+        model_outputs = sngan.train_step(
+            data_input, optim_dict, running_status=dict(iteration=1))
+        assert 'loss_disc_fake' in model_outputs['log_vars']
+        assert 'loss_disc_fake_g' in model_outputs['log_vars']
+
+        # test customized config
+        sagan = BasicConditionalGAN(
+            self.generator_cfg,
+            self.disc_cfg,
+            self.gan_loss,
+            self.disc_auxiliary_loss,
+        )
+        # test train step
+        data = torch.randn((2, 3, 32, 32))
+        data_input = dict(img=data, gt_label=label)
+        optimizer_g = torch.optim.SGD(sngan.generator.parameters(), lr=0.01)
+        optimizer_d = torch.optim.SGD(
+            sngan.discriminator.parameters(), lr=0.01)
+        optim_dict = dict(generator=optimizer_g, discriminator=optimizer_d)
+
+        model_outputs = sagan.train_step(data_input, optim_dict)
+        assert 'results' in model_outputs
+        assert 'log_vars' in model_outputs
+        assert model_outputs['num_samples'] == 2
+
+        sagan = BasicConditionalGAN(
+            self.generator_cfg, self.disc_cfg, self.gan_loss,
+            dict(
+                type='DiscShiftLoss',
+                loss_weight=0.5,
+                data_info=dict(pred='disc_pred_fake')),
+            dict(type='GeneratorPathRegularizer'))
+        assert isinstance(sagan.disc_auxiliary_losses, nn.ModuleList)
+        assert isinstance(sagan.gen_auxiliary_losses, nn.ModuleList)
+
+        sagan = BasicConditionalGAN(
+            self.generator_cfg, self.disc_cfg, self.gan_loss,
+            dict(
+                type='DiscShiftLoss',
+                loss_weight=0.5,
+                data_info=dict(pred='disc_pred_fake')),
+            [dict(type='GeneratorPathRegularizer')])
+        assert isinstance(sagan.disc_auxiliary_losses, nn.ModuleList)
+        assert isinstance(sagan.gen_auxiliary_losses, nn.ModuleList)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda')
+    def test_default_dcgan_model_cuda(self):
+        sngan = build_model(self.default_config).cuda()
+        assert isinstance(sngan, BasicConditionalGAN)
+        assert not sngan.with_disc_auxiliary_loss
+        assert sngan.with_disc
+
+        # test forward train
+        with pytest.raises(NotImplementedError):
+            _ = sngan(None, return_loss=True)
+        # test forward test
+        imgs = sngan(None, return_loss=False, mode='sampling', num_batches=2)
+        assert imgs.shape == (2, 3, 32, 32)
+
+        # test train step
+        data = torch.randn((2, 3, 32, 32)).cuda()
+        label = torch.randint(0, 10, (2, )).cuda()
+        data_input = dict(img=data, gt_label=label)
+        optimizer_g = torch.optim.SGD(sngan.generator.parameters(), lr=0.01)
+        optimizer_d = torch.optim.SGD(
+            sngan.discriminator.parameters(), lr=0.01)
+        optim_dict = dict(generator=optimizer_g, discriminator=optimizer_d)
+
+        model_outputs = sngan.train_step(data_input, optim_dict)
+        assert 'results' in model_outputs
+        assert 'log_vars' in model_outputs
+        assert model_outputs['num_samples'] == 2
+
+        # more tests for different configs with heavy computation
+        # test disc_steps
+        config_ = deepcopy(self.default_config)
+        config_['train_cfg'] = dict(disc_steps=2)
+        sngan = build_model(config_).cuda()
+        model_outputs = sngan.train_step(data_input, optim_dict)
+        assert 'loss_disc_fake' in model_outputs['log_vars']
+        assert 'loss_disc_fake_g' not in model_outputs['log_vars']
+        assert sngan.disc_steps == 2
+
+        model_outputs = sngan.train_step(
+            data_input, optim_dict, running_status=dict(iteration=1))
+        assert 'loss_disc_fake' in model_outputs['log_vars']
+        assert 'loss_disc_fake_g' in model_outputs['log_vars']

tests/test_models/test_cyclegan.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import copy
+
+import torch
+from mmcv.runner import obj_from_dict
+
+from mmgen.models import (GANLoss, L1Loss, PatchDiscriminator, ResnetGenerator,
+                          build_model)
+
+
+def test_cyclegan():
+
+    model_cfg = dict(
+        type='CycleGAN',
+        default_domain='photo',
+        reachable_domains=['photo', 'mask'],
+        related_domains=['photo', 'mask'],
+        generator=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)),
+        discriminator=dict(
+            type='PatchDiscriminator',
+            in_channels=3,
+            base_channels=64,
+            num_conv=3,
+            norm_cfg=dict(type='IN'),
+            init_cfg=dict(type='normal', gain=0.02)),
+        gan_loss=dict(
+            type='GANLoss',
+            gan_type='lsgan',
+            real_label_val=1.0,
+            fake_label_val=0.0,
+            loss_weight=1.0),
+        gen_auxiliary_loss=[
+            dict(
+                type='L1Loss',
+                loss_weight=10.0,
+                data_info=dict(pred='cycle_photo', target='real_photo'),
+                reduction='mean'),
+            dict(
+                type='L1Loss',
+                loss_weight=10.0,
+                data_info=dict(
+                    pred='cycle_mask',
+                    target='real_mask',
+                ),
+                reduction='mean'),
+            dict(
+                type='L1Loss',
+                loss_weight=0.5,
+                data_info=dict(pred='identity_photo', target='real_photo'),
+                reduction='mean'),
+            dict(
+                type='L1Loss',
+                loss_weight=0.5,
+                data_info=dict(pred='identity_mask', target='real_mask'),
+                reduction='mean')
+        ])
+
+    train_cfg = None
+    test_cfg = None
+
+    # build synthesizer
+    synthesizer = build_model(
+        model_cfg, train_cfg=train_cfg, test_cfg=test_cfg)
+
+    # test attributes
+    assert synthesizer.__class__.__name__ == 'CycleGAN'
+    assert isinstance(synthesizer.generators['photo'], ResnetGenerator)
+    assert isinstance(synthesizer.generators['mask'], ResnetGenerator)
+    assert isinstance(synthesizer.discriminators['photo'], PatchDiscriminator)
+    assert isinstance(synthesizer.discriminators['mask'], PatchDiscriminator)
+    assert isinstance(synthesizer.gan_loss, GANLoss)
+    for loss_module in synthesizer.gen_auxiliary_losses:
+        assert isinstance(loss_module, L1Loss)
+
+    # prepare data
+    inputs = torch.rand(1, 3, 64, 64)
+    targets = torch.rand(1, 3, 64, 64)
+    data_batch = {'img_mask': inputs, 'img_photo': targets}
+
+    # prepare optimizer
+    optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.5, 0.999))
+    optimizer = {
+        'generators':
+        obj_from_dict(
+            optim_cfg, torch.optim,
+            dict(params=getattr(synthesizer, 'generators').parameters())),
+        'discriminators':
+        obj_from_dict(
+            optim_cfg, torch.optim,
+            dict(params=getattr(synthesizer, 'discriminators').parameters()))
+    }
+
+    # test forward_test
+    with torch.no_grad():
+        outputs = synthesizer(inputs, target_domain='photo', test_mode=True)
+    assert torch.equal(outputs['source'], data_batch['img_mask'])
+    assert torch.is_tensor(outputs['target'])
+    assert outputs['target'].size() == (1, 3, 64, 64)
+
+    with torch.no_grad():
+        outputs = synthesizer(targets, target_domain='mask', test_mode=True)
+    assert torch.equal(outputs['source'], data_batch['img_photo'])
+    assert torch.is_tensor(outputs['target'])
+    assert outputs['target'].size() == (1, 3, 64, 64)
+
+    # test forward_train
+    with torch.no_grad():
+        outputs = synthesizer(inputs, target_domain='photo', test_mode=True)
+    assert torch.equal(outputs['source'], data_batch['img_mask'])
+    assert torch.is_tensor(outputs['target'])
+    assert outputs['target'].size() == (1, 3, 64, 64)
+
+    with torch.no_grad():
+        outputs = synthesizer(targets, target_domain='mask', test_mode=True)
+    assert torch.equal(outputs['source'], data_batch['img_photo'])
+    assert torch.is_tensor(outputs['target'])
+    assert outputs['target'].size() == (1, 3, 64, 64)
+
+    # 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_mask', 'loss_gan_d_photo', 'loss_gan_g_mask',
+            'loss_gan_g_photo', 'loss_l1'
+    ]:
+        assert isinstance(outputs['log_vars'][v], float)
+    assert outputs['num_samples'] == 1
+    assert torch.equal(outputs['results']['real_photo'],
+                       data_batch['img_photo'])
+    assert torch.equal(outputs['results']['real_mask'], data_batch['img_mask'])
+    assert torch.is_tensor(outputs['results']['fake_mask'])
+    assert torch.is_tensor(outputs['results']['fake_photo'])
+    assert outputs['results']['fake_mask'].size() == (1, 3, 64, 64)
+    assert outputs['results']['fake_photo'].size() == (1, 3, 64, 64)
+
+    # test train_step and forward_test (gpu)
+    if torch.cuda.is_available():
+        synthesizer = synthesizer.cuda()
+        optimizer = {
+            'generators':
+            obj_from_dict(
+                optim_cfg, torch.optim,
+                dict(params=getattr(synthesizer, 'generators').parameters())),
+            'discriminators':
+            obj_from_dict(
+                optim_cfg, torch.optim,
+                dict(
+                    params=getattr(synthesizer,
+                                   'discriminators').parameters()))
+        }
+        data_batch_cuda = copy.deepcopy(data_batch)
+        data_batch_cuda['img_mask'] = inputs.cuda()
+        data_batch_cuda['img_photo'] = targets.cuda()
+
+        # forward_test
+        with torch.no_grad():
+            outputs = synthesizer(
+                data_batch_cuda['img_mask'],
+                target_domain='photo',
+                test_mode=True)
+        assert torch.equal(outputs['source'],
+                           data_batch_cuda['img_mask'].cpu())
+        assert torch.is_tensor(outputs['target'])
+        assert outputs['target'].size() == (1, 3, 64, 64)
+
+        with torch.no_grad():
+            outputs = synthesizer(
+                data_batch_cuda['img_photo'],
+                target_domain='mask',
+                test_mode=True)
+        assert torch.equal(outputs['source'],
+                           data_batch_cuda['img_photo'].cpu())
+        assert torch.is_tensor(outputs['target'])
+        assert outputs['target'].size() == (1, 3, 64, 64)
+
+        # test forward_train
+        with torch.no_grad():
+            outputs = synthesizer(
+                data_batch_cuda['img_mask'],
+                target_domain='photo',
+                test_mode=False)
+        assert torch.equal(outputs['source'], data_batch_cuda['img_mask'])
+        assert torch.is_tensor(outputs['target'])
+        assert outputs['target'].size() == (1, 3, 64, 64)
+
+        with torch.no_grad():
+            outputs = synthesizer(
+                data_batch_cuda['img_photo'],
+                target_domain='mask',
+                test_mode=False)
+        assert torch.equal(outputs['source'], data_batch_cuda['img_photo'])
+        assert torch.is_tensor(outputs['target'])
+        assert outputs['target'].size() == (1, 3, 64, 64)
+
+        # train_step
+        outputs = synthesizer.train_step(data_batch_cuda, optimizer)
+        assert isinstance(outputs, dict)
+        assert isinstance(outputs['log_vars'], dict)
+        print(outputs['log_vars'].keys())
+        assert isinstance(outputs['results'], dict)
+        for v in [
+                'loss_gan_d_mask', 'loss_gan_d_photo', 'loss_gan_g_mask',
+                'loss_gan_g_photo', 'loss_l1'
+        ]:
+            assert isinstance(outputs['log_vars'][v], float)
+        assert outputs['num_samples'] == 1
+        assert torch.equal(outputs['results']['real_photo'],
+                           data_batch_cuda['img_photo'].cpu())
+        assert torch.equal(outputs['results']['real_mask'],
+                           data_batch_cuda['img_mask'].cpu())
+        assert torch.is_tensor(outputs['results']['fake_mask'])
+        assert torch.is_tensor(outputs['results']['fake_photo'])
+        assert outputs['results']['fake_mask'].size() == (1, 3, 64, 64)
+        assert outputs['results']['fake_photo'].size() == (1, 3, 64, 64)
+
+    # test disc_steps and disc_init_steps
+    data_batch['img_mask'] = inputs.cpu()
+    data_batch['img_photo'] = 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 = {
+        'generators':
+        obj_from_dict(
+            optim_cfg, torch.optim,
+            dict(params=getattr(synthesizer, 'generators').parameters())),
+        'discriminators':
+        obj_from_dict(
+            optim_cfg, torch.optim,
+            dict(params=getattr(synthesizer, 'discriminators').parameters()))
+    }
+
+    # iter 0, 1
+    for i in range(2):
+        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_g_mask', 'loss_gan_g_photo', 'loss_l1']:
+            assert outputs['log_vars'].get(v) is None
+        assert isinstance(outputs['log_vars']['loss_gan_d_mask'], float)
+        assert isinstance(outputs['log_vars']['loss_gan_d_photo'], float)
+        assert outputs['num_samples'] == 1
+        assert torch.equal(outputs['results']['real_photo'],
+                           data_batch['img_photo'])
+        assert torch.equal(outputs['results']['real_mask'],
+                           data_batch['img_mask'])
+        assert torch.is_tensor(outputs['results']['fake_mask'])
+        assert torch.is_tensor(outputs['results']['fake_photo'])
+        assert outputs['results']['fake_mask'].size() == (1, 3, 64, 64)
+        assert outputs['results']['fake_photo'].size() == (1, 3, 64, 64)
+        assert synthesizer.iteration == i + 1
+
+    # iter 2, 3, 4, 5
+    for i in range(2, 6):
+        assert synthesizer.iteration == 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_mask', 'loss_gan_d_photo', 'loss_gan_g_mask',
+            'loss_gan_g_photo', 'loss_l1'
+        ]
+        if i % 2 == 1:
+            log_None_list = ['loss_gan_g_mask', 'loss_gan_g_photo', 'loss_l1']
+            for v in log_None_list:
+                assert outputs['log_vars'].get(v) is None
+                log_check_list.remove(v)
+        for v in log_check_list:
+            assert isinstance(outputs['log_vars'][v], float)
+        assert outputs['num_samples'] == 1
+        assert torch.equal(outputs['results']['real_mask'],
+                           data_batch['img_mask'])
+        assert torch.equal(outputs['results']['real_photo'],
+                           data_batch['img_photo'])
+        assert torch.is_tensor(outputs['results']['fake_mask'])
+        assert torch.is_tensor(outputs['results']['fake_photo'])
+        assert outputs['results']['fake_mask'].size() == (1, 3, 64, 64)
+        assert outputs['results']['fake_photo'].size() == (1, 3, 64, 64)
+        assert synthesizer.iteration == i + 1
+
+    # test GAN image buffer size = 0
+    data_batch['img_mask'] = inputs.cpu()
+    data_batch['img_photo'] = targets.cpu()
+    train_cfg = dict(buffer_size=0)
+    synthesizer = build_model(
+        model_cfg, train_cfg=train_cfg, test_cfg=test_cfg)
+    optimizer = {
+        'generators':
+        obj_from_dict(
+            optim_cfg, torch.optim,
+            dict(params=getattr(synthesizer, 'generators').parameters())),
+        'discriminators':
+        obj_from_dict(
+            optim_cfg, torch.optim,
+            dict(params=getattr(synthesizer, 'discriminators').parameters()))
+    }
+    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_mask', 'loss_gan_d_photo', 'loss_gan_g_mask',
+            'loss_gan_g_photo', 'loss_l1'
+    ]:
+        assert isinstance(outputs['log_vars'][v], float)
+    assert outputs['num_samples'] == 1
+    assert torch.equal(outputs['results']['real_mask'], data_batch['img_mask'])
+    assert torch.equal(outputs['results']['real_photo'],
+                       data_batch['img_photo'])
+    assert torch.is_tensor(outputs['results']['fake_mask'])
+    assert torch.is_tensor(outputs['results']['fake_photo'])
+    assert outputs['results']['fake_mask'].size() == (1, 3, 64, 64)
+    assert outputs['results']['fake_photo'].size() == (1, 3, 64, 64)
+    assert synthesizer.iteration == 1

tests/test_models/test_mspie_styelgan2.py

+# Copyright (c) OpenMMLab. All rights reserved.
+from copy import deepcopy
+
+import torch
+
+from mmgen.models.gans.mspie_stylegan2 import MSPIEStyleGAN2
+
+
+class TestMSStyleGAN2:
+
+    @classmethod
+    def setup_class(cls):
+        cls.generator_cfg = dict(
+            type='MSStyleGANv2Generator', out_size=32, style_channels=16)
+        cls.disc_cfg = dict(
+            type='MSStyleGAN2Discriminator',
+            in_size=32,
+            with_adaptive_pool=True)
+        cls.gan_loss = dict(type='GANLoss', gan_type='vanilla')
+        cls.disc_auxiliary_loss = dict(
+            type='R1GradientPenalty',
+            loss_weight=10. / 2.,
+            interval=1,
+            norm_mode='HWC',
+            data_info=dict(real_data='real_imgs', discriminator='disc'))
+
+        cls.train_cfg = dict(
+            use_ema=True,
+            num_upblocks=3,
+            multi_input_scales=[0, 2, 4],
+            multi_scale_probability=[0.5, 0.25, 0.25])
+
+    def test_msstylegan2_cpu(self):
+        stylegan2 = MSPIEStyleGAN2(
+            self.generator_cfg,
+            self.disc_cfg,
+            self.gan_loss,
+            self.disc_auxiliary_loss,
+            None,
+            train_cfg=self.train_cfg,
+            test_cfg=None)
+
+        optimizer_g = torch.optim.SGD(
+            stylegan2.generator.parameters(), lr=0.01)
+        optimizer_d = torch.optim.SGD(
+            stylegan2.discriminator.parameters(), lr=0.01)
+        optim_dict = dict(generator=optimizer_g, discriminator=optimizer_d)
+
+        data = torch.randn((2, 3, 16, 16))
+        data_input = dict(real_img=data)
+
+        model_outputs = stylegan2.train_step(data_input, optim_dict)
+        assert 'results' in model_outputs
+        assert 'log_vars' in model_outputs
+        assert model_outputs['num_samples'] == 2
+
+        cfg_ = deepcopy(self.train_cfg)
+        cfg_['disc_steps'] = 2
+
+        stylegan2 = MSPIEStyleGAN2(
+            self.generator_cfg,
+            self.disc_cfg,
+            self.gan_loss,
+            self.disc_auxiliary_loss,
+            None,
+            train_cfg=cfg_,
+            test_cfg=None)
+
+        optimizer_g = torch.optim.SGD(
+            stylegan2.generator.parameters(), lr=0.01)
+        optimizer_d = torch.optim.SGD(
+            stylegan2.discriminator.parameters(), lr=0.01)
+        optim_dict = dict(generator=optimizer_g, discriminator=optimizer_d)
+
+        data = torch.randn((2, 3, 16, 16))
+        data_input = dict(real_img=data)
+
+        model_outputs = stylegan2.train_step(data_input, optim_dict)
+        assert 'results' in model_outputs
+        assert 'log_vars' in model_outputs
+        assert model_outputs['num_samples'] == 2

tests/test_models/test_pggan.py

+# Copyright (c) OpenMMLab. All rights reserved.
+from copy import deepcopy
+
+import numpy as np
+import pytest
+import torch
+import torch.nn as nn
+
+from mmgen.models.gans import ProgressiveGrowingGAN
+
+
+class TestPGGAN:
+
+    @classmethod
+    def setup_class(cls):
+        cls.generator_cfg = dict(
+            type='PGGANGenerator',
+            noise_size=8,
+            out_scale=16,
+            base_channels=32,
+            max_channels=32)
+
+        cls.discriminator_cfg = dict(
+            type='PGGANDiscriminator', in_scale=16, label_size=0)
+
+        cls.gan_loss = dict(type='GANLoss', gan_type='vanilla')
+        cls.disc_auxiliary_loss = [
+            dict(
+                type='DiscShiftLoss',
+                loss_weight=0.5,
+                data_info=dict(pred='disc_pred_fake')),
+            dict(
+                type='DiscShiftLoss',
+                loss_weight=0.5,
+                data_info=dict(pred='disc_pred_real'))
+        ]
+        cls.train_cfg = dict(
+            use_ema=True,
+            nkimgs_per_scale={
+                '4': 0.004,
+                '8': 0.008,
+                '16': 0.016
+            },
+            optimizer_cfg=dict(
+                generator=dict(type='Adam', lr=0.0002, betas=(0.5, 0.999)),
+                discriminator=dict(type='Adam', lr=0.0002,
+                                   betas=(0.5, 0.999))),
+            g_lr_base=0.0001,
+            d_lr_base=0.0001,
+            g_lr_schedule={'16': 0.00005})
+
+    def test_pggan_cpu(self):
+        # test default config
+        pggan = ProgressiveGrowingGAN(
+            self.generator_cfg,
+            self.discriminator_cfg,
+            self.gan_loss,
+            disc_auxiliary_loss=self.disc_auxiliary_loss,
+            train_cfg=self.train_cfg)
+
+        data_batch = dict(real_img=torch.randn(3, 3, 16, 16))
+
+        for iter_num in range(6):
+            outputs = pggan.train_step(
+                data_batch,
+                None,
+                running_status=dict(iteration=iter_num, batch_size=3))
+            results = outputs['results']
+            if iter_num == 1:
+                assert results['fake_imgs'].shape == (3, 3, 4, 4)
+            elif iter_num == 2:
+                assert results['fake_imgs'].shape == (3, 3, 8, 8)
+                assert np.isclose(pggan._actual_nkimgs[0], 0.006, atol=1e-8)
+            elif iter_num == 3:
+                assert results['fake_imgs'].shape == (3, 3, 8, 8)
+                assert np.isclose(pggan._actual_nkimgs[0], 0.006, atol=1e-8)
+                assert np.isclose(
+                    pggan.optimizer['generator'].defaults['lr'],
+                    0.0001,
+                    atol=1e-8)
+            elif iter_num == 5:
+                assert results['fake_imgs'].shape == (3, 3, 16, 16)
+                assert np.isclose(pggan._actual_nkimgs[-1], 0.012, atol=1e-8)
+                assert np.isclose(
+                    pggan.optimizer['generator'].defaults['lr'],
+                    0.00005,
+                    atol=1e-8)
+
+        # test sample from noise
+        outputs = pggan.sample_from_noise(None, num_batches=2)
+        assert outputs.shape == (4, 3, 16, 16)
+
+        outputs = pggan.sample_from_noise(
+            None,
+            num_batches=2,
+            return_noise=True,
+            transition_weight=0.2,
+            sample_model='ema')
+        assert outputs['fake_img'].shape == (2, 3, 16, 16)
+
+        outputs = pggan.sample_from_noise(
+            None, num_batches=2, return_noise=True, sample_model='orig')
+        assert outputs['fake_img'].shape == (2, 3, 16, 16)
+
+        with pytest.raises(RuntimeError):
+            data_batch = dict(real_img=torch.randn(3, 3, 4, 32))
+            _ = pggan.train_step(
+                data_batch, None, running_status=dict(iteration=5))
+
+        # test customized config
+        train_cfg_ = deepcopy(self.train_cfg)
+        train_cfg_['use_ema'] = False
+        train_cfg_['interp_real_cfg'] = dict(
+            mode='bilinear', align_corners=False)
+        train_cfg_['interp_real_cfg'] = dict(
+            mode='bilinear', align_corners=False)
+        train_cfg_['reset_optim_for_new_scale'] = False
+        pggan = ProgressiveGrowingGAN(
+            self.generator_cfg,
+            self.discriminator_cfg,
+            self.gan_loss,
+            disc_auxiliary_loss=None,
+            train_cfg=train_cfg_)
+
+        data_batch = dict(real_img=torch.randn(3, 3, 16, 16))
+
+        outputs = pggan.train_step(
+            data_batch, None, running_status=dict(iteration=0, batch_size=3))
+        results = outputs['results']
+        assert results['fake_imgs'].shape == (3, 3, 4, 4)
+        assert not pggan.with_gen_auxiliary_loss
+        assert not pggan.with_disc_auxiliary_loss
+        assert not pggan.use_ema
+
+        data_batch = dict(real_img=torch.randn(3, 3, 16, 16))
+
+        for iter_num in range(1, 3):
+            outputs = pggan.train_step(
+                data_batch,
+                None,
+                running_status=dict(iteration=iter_num, batch_size=3))
+            results = outputs['results']
+            if iter_num == 1:
+                assert results['fake_imgs'].shape == (3, 3, 4, 4)
+            elif iter_num == 2:
+                assert results['fake_imgs'].shape == (3, 3, 8, 8)
+                assert np.isclose(pggan._actual_nkimgs[0], 0.006, atol=1e-8)
+
+        train_cfg_ = deepcopy(self.train_cfg)
+        train_cfg_['optimizer_cfg'] = dict(
+            generator=dict(type='Adam', lr=0.0002, betas=(0.5, 0.999)))
+        pggan = ProgressiveGrowingGAN(
+            self.generator_cfg,
+            None,
+            None,
+            disc_auxiliary_loss=dict(
+                type='DiscShiftLoss',
+                loss_weight=0.5,
+                data_info=dict(pred='disc_pred_fake')),
+            gen_auxiliary_loss=dict(type='GeneratorPathRegularizer'),
+            train_cfg=train_cfg_)
+        assert pggan.with_gen_auxiliary_loss
+        assert isinstance(pggan.disc_auxiliary_losses, nn.ModuleList)
+        assert pggan.gan_loss is None
+        assert pggan.discriminator is None
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda')
+    def test_pggan_cuda(self):
+        # test default config
+        pggan = ProgressiveGrowingGAN(
+            self.generator_cfg,
+            self.discriminator_cfg,
+            self.gan_loss,
+            disc_auxiliary_loss=self.disc_auxiliary_loss,
+            train_cfg=self.train_cfg).cuda()
+
+        data_batch = dict(real_img=torch.randn(3, 3, 32, 32).cuda())
+
+        for iter_num in range(6):
+            outputs = pggan.train_step(
+                data_batch,
+                None,
+                running_status=dict(iteration=iter_num, batch_size=3))
+            results = outputs['results']
+            if iter_num == 1:
+                assert results['fake_imgs'].shape == (3, 3, 4, 4)
+            elif iter_num == 2:
+                assert results['fake_imgs'].shape == (3, 3, 8, 8)
+                assert np.isclose(pggan._actual_nkimgs[0], 0.006, atol=1e-8)
+            elif iter_num == 3:
+                assert results['fake_imgs'].shape == (3, 3, 8, 8)
+                assert np.isclose(pggan._actual_nkimgs[0], 0.006, atol=1e-8)
+            elif iter_num == 5:
+                assert results['fake_imgs'].shape == (3, 3, 16, 16)
+                assert np.isclose(pggan._actual_nkimgs[-1], 0.012, atol=1e-8)

tests/test_models/test_pix2pix.py

+# Copyright (c) OpenMMLab. All rights reserved.
+import copy
+
+import torch
+from mmcv.runner import obj_from_dict
+
+from mmgen.models import GANLoss, L1Loss, build_model
+from mmgen.models.architectures.pix2pix import (PatchDiscriminator,
+                                                UnetGenerator)
+
+
+def test_pix2pix():
+    # model settings
+    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.0,
+            loss_weight=1.0),
+        default_domain='photo',
+        reachable_domains=['photo'],
+        related_domains=['photo', 'mask'],
+        gen_auxiliary_loss=dict(
+            type='L1Loss',
+            loss_weight=100.0,
+            data_info=dict(pred='fake_photo', target='real_photo'),
+            reduction='mean'))
+
+    train_cfg = None
+    test_cfg = None
+
+    # build synthesizer
+    synthesizer = build_model(
+        model_cfg, train_cfg=train_cfg, test_cfg=test_cfg)
+    # test attributes
+    assert synthesizer.__class__.__name__ == 'Pix2Pix'
+    assert isinstance(synthesizer.generators['photo'], UnetGenerator)
+    assert isinstance(synthesizer.discriminators['photo'], PatchDiscriminator)
+    assert isinstance(synthesizer.gan_loss, GANLoss)
+    assert isinstance(synthesizer.gen_auxiliary_losses[0], L1Loss)
+    assert synthesizer.test_cfg is None
+
+    # prepare data
+    img_mask = torch.rand(1, 3, 256, 256)
+    img_photo = torch.rand(1, 3, 256, 256)
+    data_batch = {'img_mask': img_mask, 'img_photo': img_photo}
+
+    # prepare optimizer
+    optim_cfg = dict(type='Adam', lr=2e-4, betas=(0.5, 0.999))
+    optimizer = {
+        'generators':
+        obj_from_dict(
+            optim_cfg, torch.optim,
+            dict(params=getattr(synthesizer, 'generators').parameters())),
+        'discriminators':
+        obj_from_dict(
+            optim_cfg, torch.optim,
+            dict(params=getattr(synthesizer, 'discriminators').parameters()))
+    }
+
+    # test forward_test
+    domain = 'photo'
+    with torch.no_grad():
+        outputs = synthesizer(img_mask, target_domain=domain, test_mode=True)
+    assert torch.equal(outputs['source'], data_batch['img_mask'])
+    assert torch.is_tensor(outputs['target'])
+    assert outputs['target'].size() == (1, 3, 256, 256)
+
+    # test forward_train
+    outputs = synthesizer(img_mask, target_domain=domain, test_mode=False)
+    assert torch.equal(outputs['source'], data_batch['img_mask'])
+    assert torch.is_tensor(outputs['target'])
+    assert outputs['target'].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_l1']:
+        assert isinstance(outputs['log_vars'][v], float)
+    assert outputs['num_samples'] == 1
+
+    assert torch.equal(outputs['results']['real_mask'], data_batch['img_mask'])
+    assert torch.equal(outputs['results']['real_photo'],
+                       data_batch['img_photo'])
+    assert torch.is_tensor(outputs['results']['fake_photo'])
+    assert outputs['results']['fake_photo'].size() == (1, 3, 256, 256)
+
+    # test cuda
+    if torch.cuda.is_available():
+        synthesizer = synthesizer.cuda()
+        optimizer = {
+            'generators':
+            obj_from_dict(
+                optim_cfg, torch.optim,
+                dict(params=getattr(synthesizer, 'generators').parameters())),
+            'discriminators':
+            obj_from_dict(
+                optim_cfg, torch.optim,
+                dict(
+                    params=getattr(synthesizer,
+                                   'discriminators').parameters()))
+        }
+        data_batch_cuda = copy.deepcopy(data_batch)
+        data_batch_cuda['img_mask'] = img_mask.cuda()
+        data_batch_cuda['img_photo'] = img_photo.cuda()
+
+        # forward_test
+        with torch.no_grad():
+            outputs = synthesizer(
+                data_batch_cuda['img_mask'],
+                target_domain=domain,
+                test_mode=True)
+        assert torch.equal(outputs['source'],
+                           data_batch_cuda['img_mask'].cpu())
+        assert torch.is_tensor(outputs['target'])
+        assert outputs['target'].size() == (1, 3, 256, 256)
+
+        # test forward_train
+        outputs = synthesizer(
+            data_batch_cuda['img_mask'], target_domain=domain, test_mode=False)
+        assert torch.equal(outputs['source'], data_batch_cuda['img_mask'])
+        assert torch.is_tensor(outputs['target'])
+        assert outputs['target'].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_l1'
+        ]:
+            assert isinstance(outputs['log_vars'][v], float)
+        assert outputs['num_samples'] == 1
+        assert torch.equal(outputs['results']['real_mask'],
+                           data_batch_cuda['img_mask'].cpu())
+        assert torch.equal(outputs['results']['real_photo'],
+                           data_batch_cuda['img_photo'].cpu())
+        assert torch.is_tensor(outputs['results']['fake_photo'])
+        assert outputs['results']['fake_photo'].size() == (1, 3, 256, 256)
+
+    # test disc_steps and disc_init_steps
+    data_batch['img_mask'] = img_mask.cpu()
+    data_batch['img_photo'] = img_photo.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 = {
+        'generators':
+        obj_from_dict(
+            optim_cfg, torch.optim,
+            dict(params=getattr(synthesizer, 'generators').parameters())),
+        'discriminators':
+        obj_from_dict(
+            optim_cfg, torch.optim,
+            dict(params=getattr(synthesizer, 'discriminators').parameters()))
+    }
+
+    # iter 0, 1
+    for i in range(2):
+        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_l1') 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_mask'],
+                           data_batch['img_mask'])
+        assert torch.equal(outputs['results']['real_photo'],
+                           data_batch['img_photo'])
+        assert torch.is_tensor(outputs['results']['fake_photo'])
+        assert outputs['results']['fake_photo'].size() == (1, 3, 256, 256)
+        assert synthesizer.iteration == i + 1
+
+    # iter 2, 3, 4, 5
+    for i in range(2, 6):
+        assert synthesizer.iteration == 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_l1'
+        ]
+        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_l1')
+        for v in log_check_list:
+            assert isinstance(outputs['log_vars'][v], float)
+        assert outputs['num_samples'] == 1
+        assert torch.equal(outputs['results']['real_mask'],
+                           data_batch['img_mask'])
+        assert torch.equal(outputs['results']['real_photo'],
+                           data_batch['img_photo'])
+        assert torch.is_tensor(outputs['results']['fake_photo'])
+        assert outputs['results']['fake_photo'].size() == (1, 3, 256, 256)
+        assert synthesizer.iteration == i + 1
+
+    # test without pixel loss
+    model_cfg_ = copy.deepcopy(model_cfg)
+    model_cfg_.pop('gen_auxiliary_loss')
+    synthesizer = build_model(model_cfg_, train_cfg=None, test_cfg=None)
+    optimizer = {
+        'generators':
+        obj_from_dict(
+            optim_cfg, torch.optim,
+            dict(params=getattr(synthesizer, 'generators').parameters())),
+        'discriminators':
+        obj_from_dict(
+            optim_cfg, torch.optim,
+            dict(params=getattr(synthesizer, 'discriminators').parameters()))
+    }
+    data_batch['img_mask'] = img_mask.cpu()
+    data_batch['img_photo'] = img_photo.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_mask'], data_batch['img_mask'])
+    assert torch.equal(outputs['results']['real_photo'],
+                       data_batch['img_photo'])
+    assert torch.is_tensor(outputs['results']['fake_photo'])
+    assert outputs['results']['fake_photo'].size() == (1, 3, 256, 256)