File size: 4,600 Bytes
fb5159d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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




import unittest

import numpy as np
from caffe2.proto import caffe2_pb2
from caffe2.python import core, workspace, test_util


@unittest.skipIf(not workspace.C.has_mkldnn, "Skipping as we do not have mkldnn.")
class TestMKLBasic(test_util.TestCase):
    def testSpatialBNTestingSpeed(self):

        input_channel = 10
        X = np.random.rand(1, input_channel, 100, 100).astype(np.float32) - 0.5
        scale = np.random.rand(input_channel).astype(np.float32) + 0.5
        bias = np.random.rand(input_channel).astype(np.float32) - 0.5
        mean = np.random.randn(input_channel).astype(np.float32)
        var = np.random.rand(input_channel).astype(np.float32) + 0.5

        mkl_do = core.DeviceOption(caffe2_pb2.MKLDNN)
        # Makes sure that feed works.
        workspace.FeedBlob("X", X)
        workspace.FeedBlob("scale", scale)
        workspace.FeedBlob("bias", bias)
        workspace.FeedBlob("mean", mean)
        workspace.FeedBlob("var", var)
        workspace.FeedBlob("X_mkl", X, device_option=mkl_do)
        workspace.FeedBlob("scale_mkl", scale, device_option=mkl_do)
        workspace.FeedBlob("bias_mkl", bias, device_option=mkl_do)
        workspace.FeedBlob("mean_mkl", mean, device_option=mkl_do)
        workspace.FeedBlob("var_mkl", var, device_option=mkl_do)
        net = core.Net("test")
        # Makes sure that we can run relu.
        net.SpatialBN(["X", "scale", "bias","mean","var"], "Y", order="NCHW",
            is_test=True,
            epsilon=1e-5)
        net.SpatialBN(["X_mkl", "scale_mkl", "bias_mkl","mean_mkl","var_mkl"], "Y_mkl", order="NCHW",
            is_test=True,
            epsilon=1e-5, device_option=mkl_do)

        workspace.CreateNet(net)
        workspace.RunNet(net)
        # makes sure that the results are good.
        np.testing.assert_allclose(
            workspace.FetchBlob("Y"),
            workspace.FetchBlob("Y_mkl"),
            atol=1e-2,
            rtol=1e-2)
        runtime = workspace.BenchmarkNet(net.Proto().name, 1, 100, True)

        print("FC CPU runtime {}, MKL runtime {}.".format(runtime[1], runtime[2]))

    def testSpatialBNTrainingSpeed(self):
        input_channel = 10
        X = np.random.rand(1, input_channel, 100, 100).astype(np.float32) - 0.5
        scale = np.random.rand(input_channel).astype(np.float32) + 0.5
        bias = np.random.rand(input_channel).astype(np.float32) - 0.5
        mean = np.random.randn(input_channel).astype(np.float32)
        var = np.random.rand(input_channel).astype(np.float32) + 0.5

        #mean = np.zeros(input_channel)
        #var = np.zeros(input_channel)

        mkl_do = core.DeviceOption(caffe2_pb2.MKLDNN)
        # Makes sure that feed works.
        workspace.FeedBlob("X", X)
        workspace.FeedBlob("scale", scale)
        workspace.FeedBlob("bias", bias)
        workspace.FeedBlob("mean", mean)
        workspace.FeedBlob("var", var)
        workspace.FeedBlob("X_mkl", X, device_option=mkl_do)
        workspace.FeedBlob("scale_mkl", scale, device_option=mkl_do)
        workspace.FeedBlob("bias_mkl", bias, device_option=mkl_do)
        workspace.FeedBlob("mean_mkl", mean, device_option=mkl_do)
        workspace.FeedBlob("var_mkl", var, device_option=mkl_do)
        net = core.Net("test")
        # Makes sure that we can run relu.
        net.SpatialBN(["X", "scale", "bias","mean", "var"],
            ["Y", "mean", "var", "saved_mean", "saved_var"],
            order="NCHW",
            is_test=False,
            epsilon=1e-5)
        net.SpatialBN(["X_mkl", "scale_mkl", "bias_mkl","mean_mkl","var_mkl"],
            ["Y_mkl", "mean_mkl", "var_mkl", "saved_mean_mkl", "saved_var_mkl"],
            order="NCHW",
            is_test=False,
            epsilon=1e-5,
            device_option=mkl_do)

        workspace.CreateNet(net)
        workspace.RunNet(net)

        # makes sure that the results are good.
        np.testing.assert_allclose(
            workspace.FetchBlob("Y"),
            workspace.FetchBlob("Y_mkl"),
            atol=1e-2,
            rtol=1e-2)
        np.testing.assert_allclose(
            workspace.FetchBlob("mean"),
            workspace.FetchBlob("mean_mkl"),
            atol=1e-2,
            rtol=1e-2)
        np.testing.assert_allclose(
            workspace.FetchBlob("var"),
            workspace.FetchBlob("var_mkl"),
            atol=1e-2,
            rtol=1e-2)

        runtime = workspace.BenchmarkNet(net.Proto().name, 1, 100, True)

        print("FC CPU runtime {}, MKL runtime {}.".format(runtime[1], runtime[2]))



if __name__ == '__main__':
    unittest.main()