Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_sngan_proj.py 3.17 KB
Newer Older
limm's avatar
add tests code  
limm committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91 
# Copyright (c) OpenMMLab. All rights reserved.
import pytest
import torch

from mmgen.models.gans import BasicConditionalGAN


class TestSNGAN_PROJ:

    @classmethod
    def setup_class(cls):
        cls.generator_cfg = dict(
            type='SNGANGenerator',
            output_scale=32,
            base_channels=256,
            num_classes=10)

        cls.discriminator_cfg = dict(
            type='ProjDiscriminator',
            input_scale=32,
            base_channels=128,
            num_classes=10)

        cls.disc_auxiliary_loss = None
        cls.gan_loss = dict(type='GANLoss', gan_type='hinge')
        cls.train_cfg = None

    def test_sngan_proj_cpu(self):
        # test default config
        snganproj = BasicConditionalGAN(
            self.generator_cfg,
            self.discriminator_cfg,
            self.gan_loss,
            disc_auxiliary_loss=None,
            train_cfg=self.train_cfg)

        # test sample from noise
        outputs = snganproj.sample_from_noise(None, num_batches=2)
        assert outputs.shape == (2, 3, 32, 32)

        outputs = snganproj.sample_from_noise(
            None, num_batches=2, return_noise=True, sample_model='orig')
        assert outputs['fake_img'].shape == (2, 3, 32, 32)

        # test train step
        img = torch.randn((2, 3, 32, 32))
        lab = torch.randint(0, 10, (2, ))
        data_input = dict(img=img, gt_label=lab)
        optimizer_g = torch.optim.SGD(
            snganproj.generator.parameters(), lr=0.01)
        optimizer_d = torch.optim.SGD(
            snganproj.discriminator.parameters(), lr=0.01)
        optim_dict = dict(generator=optimizer_g, discriminator=optimizer_d)

        model_outputs = snganproj.train_step(data_input, optim_dict)
        assert 'results' in model_outputs
        assert 'log_vars' in model_outputs
        assert model_outputs['num_samples'] == 2

    @pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda')
    def test_sngan_proj_cuda(self):
        # test default config
        snganproj = BasicConditionalGAN(
            self.generator_cfg,
            self.discriminator_cfg,
            self.gan_loss,
            disc_auxiliary_loss=self.disc_auxiliary_loss,
            train_cfg=self.train_cfg).cuda()

        # test sample from noise
        outputs = snganproj.sample_from_noise(None, num_batches=2)
        assert outputs.shape == (2, 3, 32, 32)

        outputs = snganproj.sample_from_noise(
            None, num_batches=2, return_noise=True, sample_model='orig')
        assert outputs['fake_img'].shape == (2, 3, 32, 32)

        # test train step
        img = torch.randn((2, 3, 32, 32)).cuda()
        lab = torch.randint(0, 10, (2, )).cuda()
        data_input = dict(img=img, gt_label=lab)
        optimizer_g = torch.optim.SGD(
            snganproj.generator.parameters(), lr=0.01)
        optimizer_d = torch.optim.SGD(
            snganproj.discriminator.parameters(), lr=0.01)
        optim_dict = dict(generator=optimizer_g, discriminator=optimizer_d)

        model_outputs = snganproj.train_step(data_input, optim_dict)
        assert 'results' in model_outputs
        assert 'log_vars' in model_outputs
        assert model_outputs['num_samples'] == 2