Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_bias_act.py 5.27 KB
Newer Older
Yifei Yang's avatar
[Feature] Add Ops of StyleGAN3 (#2290)  
Yifei Yang 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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144 
# Copyright (c) OpenMMLab. All rights reserved.
import pytest
import torch

from mmcv.ops import bias_act
from mmcv.ops.bias_act import EasyDict

_USING_PARROTS = True
try:
    from parrots.autograd import gradcheck
except ImportError:
    from torch.autograd import gradcheck, gradgradcheck
    _USING_PARROTS = False


class TestBiasAct:

    @classmethod
    def setup_class(cls):
        cls.input_tensor = torch.randn((1, 3), requires_grad=True)
        cls.bias = torch.randn(3, requires_grad=True)

    def test_bias_act_cpu(self):
        out = bias_act(self.input_tensor, self.bias)
        assert out.shape == (1, 3)

        # test with different dim
        input_tensor = torch.randn((1, 1, 3), requires_grad=True)
        bias = torch.randn(3, requires_grad=True)
        out = bias_act(input_tensor, bias, dim=2)
        assert out.shape == (1, 1, 3)

        # test with different act
        out = bias_act(self.input_tensor, self.bias, act='relu')
        assert out.shape == (1, 3)
        out = bias_act(self.input_tensor, self.bias, act='lrelu')
        assert out.shape == (1, 3)
        out = bias_act(self.input_tensor, self.bias, act='tanh')
        assert out.shape == (1, 3)
        out = bias_act(self.input_tensor, self.bias, act='sigmoid')
        assert out.shape == (1, 3)
        out = bias_act(self.input_tensor, self.bias, act='elu')
        assert out.shape == (1, 3)
        out = bias_act(self.input_tensor, self.bias, act='selu')
        assert out.shape == (1, 3)
        out = bias_act(self.input_tensor, self.bias, act='softplus')
        assert out.shape == (1, 3)
        out = bias_act(self.input_tensor, self.bias, act='swish')
        assert out.shape == (1, 3)

        # test with different alpha
        out = bias_act(self.input_tensor, self.bias, act='lrelu', alpha=0.1)
        assert out.shape == (1, 3)

        # test with different gain
        out1 = bias_act(self.input_tensor, self.bias, act='lrelu', gain=0.2)
        out2 = bias_act(self.input_tensor, self.bias, act='lrelu', gain=0.1)
        assert torch.allclose(out1, out2 * 2)

        # test with different clamp
        out1 = bias_act(self.input_tensor, self.bias, act='lrelu', clamp=0.5)
        out2 = bias_act(self.input_tensor, self.bias, act='lrelu', clamp=0.2)
        assert out1.max() <= 0.5
        assert out2.max() <= 0.5

    @pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda')
    def test_bias_act_cuda(self):
        if _USING_PARROTS:
            gradcheck(
                bias_act, (self.input_tensor.cuda(), self.bias.cuda()),
                delta=1e-4,
                pt_atol=1e-3)
        else:
            gradcheck(
                bias_act, (self.input_tensor.cuda(), self.bias.cuda()),
                eps=1e-4,
                atol=1e-3)

            gradgradcheck(
                bias_act, (self.input_tensor.cuda(), self.bias.cuda()),
                eps=1e-4,
                atol=1e-3)

        out = bias_act(self.input_tensor.cuda(), self.bias.cuda())
        assert out.shape == (1, 3)

        # test with different dim
        input_tensor = torch.randn((1, 1, 3), requires_grad=True).cuda()
        bias = torch.randn(3, requires_grad=True).cuda()
        out = bias_act(input_tensor, bias, dim=2)
        assert out.shape == (1, 1, 3)

        # test with different act
        out = bias_act(self.input_tensor.cuda(), self.bias.cuda(), act='relu')
        assert out.shape == (1, 3)

        out = bias_act(self.input_tensor.cuda(), self.bias.cuda(), act='lrelu')
        assert out.shape == (1, 3)
        out = bias_act(self.input_tensor.cuda(), self.bias.cuda(), act='tanh')
        assert out.shape == (1, 3)
        out = bias_act(
            self.input_tensor.cuda(), self.bias.cuda(), act='sigmoid')
        assert out.shape == (1, 3)
        out = bias_act(self.input_tensor.cuda(), self.bias.cuda(), act='elu')
        assert out.shape == (1, 3)
        out = bias_act(self.input_tensor.cuda(), self.bias.cuda(), act='selu')
        assert out.shape == (1, 3)
        out = bias_act(
            self.input_tensor.cuda(), self.bias.cuda(), act='softplus')
        assert out.shape == (1, 3)
        out = bias_act(self.input_tensor.cuda(), self.bias.cuda(), act='swish')
        assert out.shape == (1, 3)

        # test with different alpha
        out = bias_act(
            self.input_tensor.cuda(), self.bias.cuda(), act='lrelu', alpha=0.1)
        assert out.shape == (1, 3)

        # test with different gain
        out1 = bias_act(
            self.input_tensor.cuda(), self.bias.cuda(), act='lrelu', gain=0.2)
        out2 = bias_act(
            self.input_tensor.cuda(), self.bias.cuda(), act='lrelu', gain=0.1)
        assert torch.allclose(out1, out2 * 2)

        # test with different clamp
        out1 = bias_act(
            self.input_tensor.cuda(), self.bias.cuda(), act='lrelu', clamp=0.5)
        out2 = bias_act(
            self.input_tensor.cuda(), self.bias.cuda(), act='lrelu', clamp=0.2)
        assert out1.max() <= 0.5
        assert out2.max() <= 0.5

    def test_easy_dict(self):
        easy_dict = EasyDict(
            func=lambda x, **_: x,
            def_alpha=0,
            def_gain=1,
            cuda_idx=1,
            ref='',
            has_2nd_grad=False)
        _ = easy_dict.def_alpha
        easy_dict.def_alpha = 1
        del easy_dict.def_alpha