Merge pull request #255 from salexspb/caffe2_py_test

caffe2 test fix + validation

Author: prigoyal

test_python/test_c2.py

@@ -15,33 +15,55 @@
 
 import unittest, os
 import numpy as np
+import hypothesis.strategies as st
+import caffe2.python.hypothesis_test_util as hu
 
-from caffe2.proto import caffe2_pb2
-from caffe2.python import core, workspace, dyndep
+from hypothesis import given
+from caffe2.python import core, dyndep
 
-tc_c2_lib = os.path.join(os.environ.get("CONDA_PREFIX"), "lib/libtc_c2.so")
-dyndep.InitOpsLibrary(tc_c2_lib)
 
+CONDA_PREFIX = os.environ.get("CONDA_PREFIX")
+if CONDA_PREFIX:
+    tc_c2_lib = os.path.join(CONDA_PREFIX, "lib/libtc_c2.so")
+else:
+    dyndep.InitOpsLibrary("@/tc/tc:tc_c2")
 
-class TestCaffe2(unittest.TestCase):
 
-    def test_matmul_caffe2(self):
-        lang = """
+class TestCaffe2(hu.HypothesisTestCase):
+    @given(n=st.integers(1, 128),
+           m=st.integers(1, 128),
+           k=st.integers(1, 128),
+           seed=st.integers(min_value=0, max_value=2**32 - 1),
+           **hu.gcs_gpu_only)
+    def test_matmul(self, n, m, k, seed, gc, dc):
+        np.random.seed(seed)
+
+        tc_forward = """
         def matmul(float(M,N) A, float(N,K) B) -> (output) {
           output(i, j) +=! A(i, kk) * B(kk, j)
         }
         """
+
         # TODO: (prigoyal) serialize the options
         # options = Options("mlp")
-        mat1, mat2 = np.random.rand(100, 400), np.random.rand(400, 500)
-        with core.DeviceScope(core.DeviceOption(caffe2_pb2.CUDA, 0)):
-            workspace.FeedBlob('mat1', mat1.astype(np.float32))
-            workspace.FeedBlob('mat2', mat2.astype(np.float32))
-            matmul = core.CreateOperator(
-                "TcOp", ["mat1", "mat2"], ["out"], lang=lang, tcName="matmul"
-            )
-        workspace.RunOperatorOnce(matmul)
-        out = workspace.FetchBlob("out")
+        X = np.random.rand(m, k).astype(np.float32)
+        W = np.random.rand(k, n).astype(np.float32)
+
+        def ref(X, W):
+            return [np.dot(X, W)]
+
+        op = core.CreateOperator(
+            "TcOp", ["X", "Y"], "out",
+            tcDef=tc_forward,
+            tcName="matmul",
+        )
+
+        self.assertReferenceChecks(
+            device_option=gc,
+            op=op,
+            inputs=[X, W],
+            reference=ref,
+        )
 
 
 if __name__ == '__main__':