File: mkl_packed_fc_op_test.py

package info (click to toggle)
pytorch 1.13.1%2Bdfsg-4
  • links: PTS, VCS
  • area: main
  • in suites: bookworm
  • size: 139,252 kB
  • sloc: cpp: 1,100,274; python: 706,454; ansic: 83,052; asm: 7,618; java: 3,273; sh: 2,841; javascript: 612; makefile: 323; xml: 269; ruby: 185; yacc: 144; objc: 68; lex: 44
file content (76 lines) | stat: -rw-r--r-- 2,647 bytes parent folder | download | duplicates (2)
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





import unittest
import hypothesis.strategies as st
from hypothesis import given
import numpy as np
from caffe2.python import core
import caffe2.python.hypothesis_test_util as hu


@unittest.skipIf(not core.IsOperator("PackedFC"),
                 "PackedFC is not supported in this caffe2 build.")
class PackedFCTest(hu.HypothesisTestCase):
    @given(seed=st.integers(0, 65536),
           M=st.integers(16, 32),
           K=st.integers(128, 1024),
           N=st.integers(128, 1024),
           **hu.gcs_cpu_only)
    @unittest.skipIf(not core.C.builtin_cpu_supports_avx2(),
                     "Intel MKL sgemm_pack has a known numerical issue with "
                     "non-avx2 machines that will be fixed in a later build.")
    def test_packed_fc(self, seed, M, K, N, gc, dc):
        np.random.seed(seed)
        X = np.random.rand(M, K).astype(np.float32) - 0.5
        W = np.random.rand(N, K).astype(np.float32) - 0.5
        b = np.random.rand(N).astype(np.float32) - 0.5

        # If you are debugging, the following hard-coded ones might help.
        # X = np.ones((24, 256)).astype(np.float32)
        # W = np.ones((128, 256)).astype(np.float32)
        # b = np.zeros(128).astype(np.float32)

        def ref(X, W, b):
            return (np.dot(X, W.T) + b,)

        for name in ["FC", "PackedFC"]:
            op = core.CreateOperator(
                name,
                ["X", "W", "b"],
                ["Y"],
            )
            self.assertReferenceChecks(gc, op, [X, W, b], ref)

    @unittest.skipIf(not core.C.builtin_cpu_supports_avx2(),
                     "Intel MKL sgemm_pack has a known numerical issue with "
                     "non-avx2 machines that will be fixed in a later build.")
    @given(axis=st.integers(min_value=1, max_value=4),
           num_output=st.integers(min_value=4, max_value=8),
           **hu.gcs_cpu_only)
    def test_packed_fc_axis(self, axis, num_output, gc, dc):
        np.random.seed(1701)
        X = np.random.randn(1, 2, 3, 2, 1).astype(np.float32)
        K = np.prod(X.shape[axis:])
        N = num_output
        W = np.random.randn(N, K).astype(np.float32)
        b = np.random.randn(N).astype(np.float32)

        op = core.CreateOperator(
            "PackedFC",
            ["X", "W", "b"],
            ["Y"],
            axis=axis)

        def ref(X, W, b):
            output_axes = list(X.shape[:axis]) + [N]
            return (
                np.dot(X.reshape(int(X.size / K), K), W.T).reshape(output_axes) + b,)

        self.assertReferenceChecks(gc, op, [X, W, b], ref)

if __name__ == "__main__":
    import unittest
    unittest.main()