pytorch

1
import numpy as np
2
import unittest
3

4
import caffe2.python.fakelowp.init_shared_libs  # noqa
5
from hypothesis import given, settings
6
from hypothesis import strategies as st
7
from caffe2.proto import caffe2_pb2
8
from caffe2.python import core
9
from caffe2.python import workspace
10
from caffe2.python.onnx.onnxifi import onnxifi_caffe2_net
11
from caffe2.python.fakelowp.test_utils import print_test_debug_info
12
import caffe2.python.serialized_test.serialized_test_util as serial
13
import datetime
14

15
core.GlobalInit(["caffe2", "--glow_global_fp16=1",
16
                 "--glow_global_fused_scale_offset_fp16=1",
17
                 "--glow_global_force_sls_fp16_accum=1"])
18

19
GLOW_LOWERED_BATCHNORM = False
20

21

22
def reference_spatialbn_test16(X, scale, bias, mean, var, epsilon, order):
23
    X = X.astype(np.float16)
24
    scale = scale.astype(np.float16)
25
    bias = bias.astype(np.float16)
26
    mean = mean.astype(np.float16)
27
    # var = var.astype(np.float16)
28
    assert(order == "NCHW")
29

30
    scale = scale[np.newaxis, :, np.newaxis, np.newaxis]
31
    bias = bias[np.newaxis, :, np.newaxis, np.newaxis]
32
    mean = mean[np.newaxis, :, np.newaxis, np.newaxis]
33
    var = var[np.newaxis, :, np.newaxis, np.newaxis]
34
    Y = ((X - mean) * (scale / np.sqrt(var + epsilon).astype(np.float16))) + bias
35
    return Y.astype(np.float32)
36

37

38
# Test the lowered BN op
39
class BatchnormTest(serial.SerializedTestCase):
40
    # TODO: using hypothesis seed, sweep dimensions
41
    @given(seed=st.integers(0, 65535),
42
           size=st.integers(2, 30),
43
           input_channels=st.integers(2, 40),
44
           batch_size=st.integers(2, 20))
45
    @settings(deadline=datetime.timedelta(seconds=10))
46
    def test_bn(self, seed, size, input_channels, batch_size):
47
        workspace.ResetWorkspace()
48
        np.random.seed(seed)
49

50
        order = "NCHW"
51
        epsilon = 1e-3
52

53
        pred_net = caffe2_pb2.NetDef()
54
        pred_net.name = "pred"
55
        pred_net.external_input.extend(["X", "scale", "bias", "mean", "var"])
56
        pred_net.external_output.append("Y")
57
        pred_net.op.add().CopyFrom(
58
            core.CreateOperator(
59
                "SpatialBN",
60
                ["X", "scale", "bias", "mean", "var"],
61
                ["Y"],
62
                order=order,
63
                is_test=True,
64
                epsilon=epsilon
65
            )
66
        )
67

68
        if GLOW_LOWERED_BATCHNORM:
69
            refopname = "SpatialBNFakeLoweredFp16NNPI"
70
        else:
71
            refopname = "SpatialBNFakeFp16NNPI"
72

73
        pred_net_ref = caffe2_pb2.NetDef()
74
        pred_net_ref.name = "pred"
75
        pred_net_ref.external_input.extend(["X", "scale", "bias", "mean", "var"])
76
        pred_net_ref.external_output.append("X")
77
        pred_net_ref.op.add().CopyFrom(
78
            core.CreateOperator(
79
                refopname,
80
                ["X", "scale", "bias", "mean", "var"],
81
                ["Y"],
82
                order=order,
83
                is_test=True,
84
                epsilon=epsilon
85
            )
86
        )
87

88
        scale = np.random.rand(input_channels).astype(np.float32) + 0.5
89
        bias = np.random.rand(input_channels).astype(np.float32) - 0.5
90
        mean = np.random.randn(input_channels).astype(np.float32)
91
        var = np.random.rand(input_channels).astype(np.float32) + 0.5
92
        X = np.random.rand(
93
            batch_size, input_channels, size, size).astype(np.float32) - 0.5
94

95
        workspace.FeedBlob("scale", scale)
96
        workspace.FeedBlob("bias", bias)
97
        workspace.FeedBlob("mean", mean)
98
        workspace.FeedBlob("var", var)
99

100
        # Use for reference to debug
101
        # Y_np = reference_spatialbn_test16(X, scale, bias, mean, var, epsilon, order)
102

103
        pred_net_onnxified = onnxifi_caffe2_net(
104
            pred_net,
105
            {"X": [batch_size, input_channels, size, size],
106
             "scale": [input_channels],
107
             "bias": [input_channels],
108
             "mean": [input_channels],
109
             "var": [input_channels]},
110
            debug=True,
111
            adjust_batch=False,
112
            use_onnx=False
113
        )
114
        num_onnxified_ops = sum(
115
            1 if o.type == "Onnxifi" else 0 for o in pred_net_onnxified.op)
116
        np.testing.assert_equal(num_onnxified_ops, 1)
117

118
        workspace.FeedBlob("X", X)
119

120
        workspace.CreateNet(pred_net_onnxified)
121
        workspace.CreateNet(pred_net_ref)
122

123
        workspace.RunNet(pred_net_ref.name)
124
        Y_c2 = workspace.FetchBlob("Y")
125

126
        workspace.RunNet(pred_net_onnxified.name)
127
        Y_glow = workspace.FetchBlob("Y")
128

129
        if not np.allclose(Y_glow.astype(np.float16), Y_c2.astype(np.float16)):
130
            diff = np.abs(Y_glow - Y_c2).astype(np.float16)
131
            print_test_debug_info(
132
                "bn",
133
                {
134
                    "seed": seed,
135
                    "scale": scale,
136
                    "bias": bias,
137
                    "mean": mean,
138
                    "var": var,
139
                    "Y_np": Y_c2,
140
                    "Y_glow": Y_glow,
141
                    "diff": diff,
142
                    "rowwise_diff": np.max(np.abs(diff), -1)})
143
            assert(0)
144