Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.

1
import unittest
2

3
import torch
4
import torch.nn as nn
5

6
from apex.fp16_utils import FP16Model
7

8

9
class DummyBlock(nn.Module):
10
    def __init__(self):
11
        super(DummyBlock, self).__init__()
12

13
        self.conv = nn.Conv2d(10, 10, 2)
14
        self.bn = nn.BatchNorm2d(10, affine=True)
15

16
    def forward(self, x):
17
        return self.conv(self.bn(x))
18

19

20
class DummyNet(nn.Module):
21
    def __init__(self):
22
        super(DummyNet, self).__init__()
23

24
        self.conv1 = nn.Conv2d(3, 10, 2)
25
        self.bn1 = nn.BatchNorm2d(10, affine=False)
26
        self.db1 = DummyBlock()
27
        self.db2 = DummyBlock()
28

29
    def forward(self, x):
30
        out = x
31
        out = self.conv1(out)
32
        out = self.bn1(out)
33
        out = self.db1(out)
34
        out = self.db2(out)
35
        return out
36

37

38
class DummyNetWrapper(nn.Module):
39
    def __init__(self):
40
        super(DummyNetWrapper, self).__init__()
41

42
        self.bn = nn.BatchNorm2d(3, affine=True)
43
        self.dn = DummyNet()
44

45
    def forward(self, x):
46
        return self.dn(self.bn(x))
47

48

49
class TestFP16Model(unittest.TestCase):
50
    def setUp(self):
51
        self.N = 64
52
        self.C_in = 3
53
        self.H_in = 16
54
        self.W_in = 32
55
        self.in_tensor = torch.randn((self.N, self.C_in, self.H_in, self.W_in)).cuda()
56
        self.orig_model = DummyNetWrapper().cuda()
57
        self.fp16_model = FP16Model(self.orig_model)
58

59
    def test_params_and_buffers(self):
60
        exempted_modules = [
61
            self.fp16_model.network.bn,
62
            self.fp16_model.network.dn.db1.bn,
63
            self.fp16_model.network.dn.db2.bn,
64
        ]
65
        for m in self.fp16_model.modules():
66
            expected_dtype = torch.float if (m in exempted_modules) else torch.half
67
            for p in m.parameters(recurse=False):
68
                assert p.dtype == expected_dtype
69
            for b in m.buffers(recurse=False):
70
                assert b.dtype in (expected_dtype, torch.int64)
71

72
    def test_output_is_half(self):
73
        out_tensor = self.fp16_model(self.in_tensor)
74
        assert out_tensor.dtype == torch.half
75

76

77