test_hsigmoid.py 1.11 KB
Newer Older
limm's avatar
limm committed
1
# Copyright (c) OpenMMLab. All rights reserved.
yamengxi's avatar
yamengxi committed
2
import pytest
3
4
5
6
7
8
import torch

from mmcv.cnn.bricks import HSigmoid


def test_hsigmoid():
yamengxi's avatar
yamengxi committed
9
10
11
12
13
    # test assertion divisor can not be zero
    with pytest.raises(AssertionError):
        HSigmoid(divisor=0)

    # test with default parameters
14
15
16
17
18
    act = HSigmoid()
    input_shape = torch.Size([1, 3, 64, 64])
    input = torch.randn(input_shape)
    output = act(input)
    expected_output = torch.min(
limm's avatar
limm committed
19
        torch.max((input + 3) / 6, torch.zeros(input_shape)),
20
21
22
23
24
        torch.ones(input_shape))
    # test output shape
    assert output.shape == expected_output.shape
    # test output value
    assert torch.equal(output, expected_output)
yamengxi's avatar
yamengxi committed
25
26

    # test with designated parameters
limm's avatar
limm committed
27
    act = HSigmoid(1, 2, 0, 1)
yamengxi's avatar
yamengxi committed
28
29
30
31
    input_shape = torch.Size([1, 3, 64, 64])
    input = torch.randn(input_shape)
    output = act(input)
    expected_output = torch.min(
limm's avatar
limm committed
32
        torch.max((input + 1) / 2, torch.zeros(input_shape)),
yamengxi's avatar
yamengxi committed
33
34
35
36
37
        torch.ones(input_shape))
    # test output shape
    assert output.shape == expected_output.shape
    # test output value
    assert torch.equal(output, expected_output)