| import errno |
| import os |
| import shutil |
| import tempfile |
| import unittest |
| from collections import namedtuple |
| from typing import List |
| |
| import caffe2.python.hypothesis_test_util as htu |
| import hypothesis.strategies as st |
| import numpy as np |
| import torch |
| from torch import Tensor |
| from caffe2.proto import caffe2_pb2 |
| from caffe2.python import core, test_util, workspace, model_helper, brew |
| from hypothesis import given, settings |
| |
| |
| class TestWorkspace(unittest.TestCase): |
| def setUp(self): |
| self.net = core.Net("test-net") |
| self.testblob_ref = self.net.ConstantFill( |
| [], "testblob", shape=[1, 2, 3, 4], value=1.0 |
| ) |
| workspace.ResetWorkspace() |
| |
| def testWorkspaceHasBlobWithNonexistingName(self): |
| self.assertEqual(workspace.HasBlob("non-existing"), False) |
| |
| def testRunOperatorOnce(self): |
| self.assertEqual( |
| workspace.RunOperatorOnce(self.net.Proto().op[0].SerializeToString()), True |
| ) |
| self.assertEqual(workspace.HasBlob("testblob"), True) |
| blobs = workspace.Blobs() |
| self.assertEqual(len(blobs), 1) |
| self.assertEqual(blobs[0], "testblob") |
| |
| def testGetOperatorCost(self): |
| op = core.CreateOperator( |
| "Conv2D", |
| ["X", "W"], |
| ["Y"], |
| stride_h=1, |
| stride_w=1, |
| pad_t=1, |
| pad_l=1, |
| pad_b=1, |
| pad_r=1, |
| kernel=3, |
| ) |
| X = np.zeros((1, 8, 8, 8)) |
| W = np.zeros((1, 1, 3, 3)) |
| workspace.FeedBlob("X", X) |
| workspace.FeedBlob("W", W) |
| op_cost = workspace.GetOperatorCost(op.SerializeToString(), ["X", "W"]) |
| self.assertTupleEqual( |
| op_cost, |
| namedtuple("Cost", ["flops", "bytes_written", "bytes_read"])( |
| 1152, 256, 4168 |
| ), |
| ) |
| |
| def testRunNetOnce(self): |
| self.assertEqual( |
| workspace.RunNetOnce(self.net.Proto().SerializeToString()), True |
| ) |
| self.assertEqual(workspace.HasBlob("testblob"), True) |
| |
| def testCurrentWorkspaceWrapper(self): |
| self.assertNotIn("testblob", workspace.C.Workspace.current.blobs) |
| self.assertEqual( |
| workspace.RunNetOnce(self.net.Proto().SerializeToString()), True |
| ) |
| self.assertEqual(workspace.HasBlob("testblob"), True) |
| self.assertIn("testblob", workspace.C.Workspace.current.blobs) |
| workspace.ResetWorkspace() |
| self.assertNotIn("testblob", workspace.C.Workspace.current.blobs) |
| |
| def testRunPlan(self): |
| plan = core.Plan("test-plan") |
| plan.AddStep(core.ExecutionStep("test-step", self.net)) |
| self.assertEqual(workspace.RunPlan(plan.Proto().SerializeToString()), True) |
| self.assertEqual(workspace.HasBlob("testblob"), True) |
| |
| def testRunPlanInBackground(self): |
| plan = core.Plan("test-plan") |
| plan.AddStep(core.ExecutionStep("test-step", self.net)) |
| background_plan = workspace.RunPlanInBackground(plan) |
| while not background_plan.is_done(): |
| pass |
| self.assertEqual(background_plan.is_succeeded(), True) |
| self.assertEqual(workspace.HasBlob("testblob"), True) |
| |
| def testConstructPlanFromSteps(self): |
| step = core.ExecutionStep("test-step-as-plan", self.net) |
| self.assertEqual(workspace.RunPlan(step), True) |
| self.assertEqual(workspace.HasBlob("testblob"), True) |
| |
| def testResetWorkspace(self): |
| self.assertEqual( |
| workspace.RunNetOnce(self.net.Proto().SerializeToString()), True |
| ) |
| self.assertEqual(workspace.HasBlob("testblob"), True) |
| self.assertEqual(workspace.ResetWorkspace(), True) |
| self.assertEqual(workspace.HasBlob("testblob"), False) |
| |
| def testTensorAccess(self): |
| ws = workspace.C.Workspace() |
| |
| """ test in-place modification """ |
| ws.create_blob("tensor").feed(np.array([1.1, 1.2, 1.3])) |
| tensor = ws.blobs["tensor"].tensor() |
| tensor.data[0] = 3.3 |
| val = np.array([3.3, 1.2, 1.3]) |
| np.testing.assert_array_equal(tensor.data, val) |
| np.testing.assert_array_equal(ws.blobs["tensor"].fetch(), val) |
| |
| """ test in-place initialization """ |
| tensor.init([2, 3], core.DataType.INT32) |
| for x in range(2): |
| for y in range(3): |
| tensor.data[x, y] = 0 |
| tensor.data[1, 1] = 100 |
| val = np.zeros([2, 3], dtype=np.int32) |
| val[1, 1] = 100 |
| np.testing.assert_array_equal(tensor.data, val) |
| np.testing.assert_array_equal(ws.blobs["tensor"].fetch(), val) |
| |
| """ strings cannot be initialized from python """ |
| with self.assertRaises(RuntimeError): |
| tensor.init([3, 4], core.DataType.STRING) |
| |
| """ feed (copy) data into tensor """ |
| val = np.array([[b"abc", b"def"], [b"ghi", b"jkl"]], dtype=np.object) |
| tensor.feed(val) |
| self.assertEquals(tensor.data[0, 0], b"abc") |
| np.testing.assert_array_equal(ws.blobs["tensor"].fetch(), val) |
| |
| val = np.array([1.1, 10.2]) |
| tensor.feed(val) |
| val[0] = 5.2 |
| self.assertEquals(tensor.data[0], 1.1) |
| |
| """ fetch (copy) data from tensor """ |
| val = np.array([1.1, 1.2]) |
| tensor.feed(val) |
| val2 = tensor.fetch() |
| tensor.data[0] = 5.2 |
| val3 = tensor.fetch() |
| np.testing.assert_array_equal(val, val2) |
| self.assertEquals(val3[0], 5.2) |
| |
| def testFetchFeedBlob(self): |
| self.assertEqual( |
| workspace.RunNetOnce(self.net.Proto().SerializeToString()), True |
| ) |
| fetched = workspace.FetchBlob("testblob") |
| # check if fetched is correct. |
| self.assertEqual(fetched.shape, (1, 2, 3, 4)) |
| np.testing.assert_array_equal(fetched, 1.0) |
| fetched[:] = 2.0 |
| self.assertEqual(workspace.FeedBlob("testblob", fetched), True) |
| fetched_again = workspace.FetchBlob("testblob") |
| self.assertEqual(fetched_again.shape, (1, 2, 3, 4)) |
| np.testing.assert_array_equal(fetched_again, 2.0) |
| |
| def testFetchFeedBlobViaBlobReference(self): |
| self.assertEqual( |
| workspace.RunNetOnce(self.net.Proto().SerializeToString()), True |
| ) |
| fetched = workspace.FetchBlob(self.testblob_ref) |
| # check if fetched is correct. |
| self.assertEqual(fetched.shape, (1, 2, 3, 4)) |
| np.testing.assert_array_equal(fetched, 1.0) |
| fetched[:] = 2.0 |
| self.assertEqual(workspace.FeedBlob(self.testblob_ref, fetched), True) |
| fetched_again = workspace.FetchBlob("testblob") # fetch by name now |
| self.assertEqual(fetched_again.shape, (1, 2, 3, 4)) |
| np.testing.assert_array_equal(fetched_again, 2.0) |
| |
| def testFetchFeedBlobTypes(self): |
| for dtype in [ |
| np.float16, |
| np.float32, |
| np.float64, |
| np.bool, |
| np.int8, |
| np.int16, |
| np.int32, |
| np.int64, |
| np.uint8, |
| np.uint16, |
| ]: |
| try: |
| rng = np.iinfo(dtype).max * 2 |
| except ValueError: |
| rng = 1000 |
| data = ((np.random.rand(2, 3, 4) - 0.5) * rng).astype(dtype) |
| self.assertEqual(workspace.FeedBlob("testblob_types", data), True) |
| fetched_back = workspace.FetchBlob("testblob_types") |
| self.assertEqual(fetched_back.shape, (2, 3, 4)) |
| self.assertEqual(fetched_back.dtype, dtype) |
| np.testing.assert_array_equal(fetched_back, data) |
| |
| def testFetchFeedBlobBool(self): |
| """Special case for bool to ensure coverage of both true and false.""" |
| data = np.zeros((2, 3, 4)).astype(np.bool) |
| data.flat[::2] = True |
| self.assertEqual(workspace.FeedBlob("testblob_types", data), True) |
| fetched_back = workspace.FetchBlob("testblob_types") |
| self.assertEqual(fetched_back.shape, (2, 3, 4)) |
| self.assertEqual(fetched_back.dtype, np.bool) |
| np.testing.assert_array_equal(fetched_back, data) |
| |
| def testGetBlobSizeBytes(self): |
| for dtype in [ |
| np.float16, |
| np.float32, |
| np.float64, |
| np.bool, |
| np.int8, |
| np.int16, |
| np.int32, |
| np.int64, |
| np.uint8, |
| np.uint16, |
| ]: |
| data = np.random.randn(2, 3).astype(dtype) |
| self.assertTrue(workspace.FeedBlob("testblob_sizeBytes", data), True) |
| self.assertEqual( |
| workspace.GetBlobSizeBytes("testblob_sizeBytes"), |
| 6 * np.dtype(dtype).itemsize, |
| ) |
| strs1 = np.array([b"Hello World!", b"abcd"]) |
| strs2 = np.array([b"element1", b"element2"]) |
| strs1_len, strs2_len = 0, 0 |
| for str in strs1: |
| strs1_len += len(str) |
| for str in strs2: |
| strs2_len += len(str) |
| self.assertTrue(workspace.FeedBlob("testblob_str1", strs1), True) |
| self.assertTrue(workspace.FeedBlob("testblob_str2", strs2), True) |
| # size of blob "testblob_str1" = size_str1 * meta_.itemsize() + strs1_len |
| # size of blob "testblob_str2" = size_str2 * meta_.itemsize() + strs2_len |
| self.assertEqual( |
| workspace.GetBlobSizeBytes("testblob_str1") |
| - workspace.GetBlobSizeBytes("testblob_str2"), |
| strs1_len - strs2_len, |
| ) |
| |
| def testFetchFeedBlobZeroDim(self): |
| data = np.empty(shape=(2, 0, 3), dtype=np.float32) |
| self.assertEqual(workspace.FeedBlob("testblob_empty", data), True) |
| fetched_back = workspace.FetchBlob("testblob_empty") |
| self.assertEqual(fetched_back.shape, (2, 0, 3)) |
| self.assertEqual(fetched_back.dtype, np.float32) |
| |
| def testFetchFeedLongStringTensor(self): |
| # long strings trigger array of object creation |
| strs = np.array( |
| [ |
| b" ".join(10 * [b"long string"]), |
| b" ".join(128 * [b"very long string"]), |
| b"small \0\1\2 string", |
| b"Hello, world! I have special \0 symbols \1!", |
| ] |
| ) |
| workspace.FeedBlob("my_str_tensor", strs) |
| strs2 = workspace.FetchBlob("my_str_tensor") |
| self.assertEqual(strs.shape, strs2.shape) |
| for i in range(0, strs.shape[0]): |
| self.assertEqual(strs[i], strs2[i]) |
| |
| def testFetchFeedShortStringTensor(self): |
| # small strings trigger NPY_STRING array |
| strs = np.array([b"elem1", b"elem 2", b"element 3"]) |
| workspace.FeedBlob("my_str_tensor_2", strs) |
| strs2 = workspace.FetchBlob("my_str_tensor_2") |
| self.assertEqual(strs.shape, strs2.shape) |
| for i in range(0, strs.shape[0]): |
| self.assertEqual(strs[i], strs2[i]) |
| |
| def testFetchFeedPlainString(self): |
| # this is actual string, not a tensor of strings |
| s = b"Hello, world! I have special \0 symbols \1!" |
| workspace.FeedBlob("my_plain_string", s) |
| s2 = workspace.FetchBlob("my_plain_string") |
| self.assertEqual(s, s2) |
| |
| def testFetchBlobs(self): |
| s1 = b"test1" |
| s2 = b"test2" |
| workspace.FeedBlob("s1", s1) |
| workspace.FeedBlob("s2", s2) |
| fetch1, fetch2 = workspace.FetchBlobs(["s1", "s2"]) |
| self.assertEquals(s1, fetch1) |
| self.assertEquals(s2, fetch2) |
| |
| def testFetchFeedViaBlobDict(self): |
| self.assertEqual( |
| workspace.RunNetOnce(self.net.Proto().SerializeToString()), True |
| ) |
| fetched = workspace.blobs["testblob"] |
| # check if fetched is correct. |
| self.assertEqual(fetched.shape, (1, 2, 3, 4)) |
| np.testing.assert_array_equal(fetched, 1.0) |
| fetched[:] = 2.0 |
| workspace.blobs["testblob"] = fetched |
| fetched_again = workspace.blobs["testblob"] |
| self.assertEqual(fetched_again.shape, (1, 2, 3, 4)) |
| np.testing.assert_array_equal(fetched_again, 2.0) |
| |
| self.assertTrue("testblob" in workspace.blobs) |
| self.assertFalse("non_existant" in workspace.blobs) |
| self.assertEqual(len(workspace.blobs), 1) |
| for key in workspace.blobs: |
| self.assertEqual(key, "testblob") |
| |
| def testTorchInterop(self): |
| workspace.RunOperatorOnce( |
| core.CreateOperator( |
| "ConstantFill", [], "foo", shape=(4,), value=2, dtype=10 |
| ) |
| ) |
| t = workspace.FetchTorch("foo") |
| t.resize_(5) |
| t[4] = t[2] = 777 |
| np.testing.assert_array_equal(t.numpy(), np.array([2, 2, 777, 2, 777])) |
| np.testing.assert_array_equal( |
| workspace.FetchBlob("foo"), np.array([2, 2, 777, 2, 777]) |
| ) |
| |
| z = torch.ones((4,), dtype=torch.int64) |
| workspace.FeedBlob("bar", z) |
| workspace.RunOperatorOnce( |
| core.CreateOperator("Reshape", ["bar"], ["bar", "_"], shape=(2, 2)) |
| ) |
| z[0, 1] = 123 |
| np.testing.assert_array_equal( |
| workspace.FetchBlob("bar"), np.array([[1, 123], [1, 1]]) |
| ) |
| np.testing.assert_array_equal(z, np.array([[1, 123], [1, 1]])) |
| |
| |
| class TestMultiWorkspaces(unittest.TestCase): |
| def setUp(self): |
| workspace.SwitchWorkspace("default") |
| workspace.ResetWorkspace() |
| |
| def testCreateWorkspace(self): |
| self.net = core.Net("test-net") |
| self.net.ConstantFill([], "testblob", shape=[1, 2, 3, 4], value=1.0) |
| self.assertEqual( |
| workspace.RunNetOnce(self.net.Proto().SerializeToString()), True |
| ) |
| self.assertEqual(workspace.HasBlob("testblob"), True) |
| self.assertEqual(workspace.SwitchWorkspace("test", True), None) |
| self.assertEqual(workspace.HasBlob("testblob"), False) |
| self.assertEqual(workspace.SwitchWorkspace("default"), None) |
| self.assertEqual(workspace.HasBlob("testblob"), True) |
| |
| try: |
| # The following should raise an error. |
| workspace.SwitchWorkspace("non-existing") |
| # so this should never happen. |
| self.assertEqual(True, False) |
| except RuntimeError: |
| pass |
| |
| workspaces = workspace.Workspaces() |
| self.assertTrue("default" in workspaces) |
| self.assertTrue("test" in workspaces) |
| |
| |
| @unittest.skipIf(not workspace.has_gpu_support, "No gpu support.") |
| class TestWorkspaceGPU(test_util.TestCase): |
| def setUp(self): |
| workspace.ResetWorkspace() |
| self.net = core.Net("test-net") |
| self.net.ConstantFill([], "testblob", shape=[1, 2, 3, 4], value=1.0) |
| self.net.RunAllOnGPU() |
| |
| def testFetchBlobGPU(self): |
| self.assertEqual( |
| workspace.RunNetOnce(self.net.Proto().SerializeToString()), True |
| ) |
| fetched = workspace.FetchBlob("testblob") |
| # check if fetched is correct. |
| self.assertEqual(fetched.shape, (1, 2, 3, 4)) |
| np.testing.assert_array_equal(fetched, 1.0) |
| fetched[:] = 2.0 |
| self.assertEqual(workspace.FeedBlob("testblob", fetched), True) |
| fetched_again = workspace.FetchBlob("testblob") |
| self.assertEqual(fetched_again.shape, (1, 2, 3, 4)) |
| np.testing.assert_array_equal(fetched_again, 2.0) |
| |
| def testGetGpuPeerAccessPattern(self): |
| pattern = workspace.GetGpuPeerAccessPattern() |
| self.assertEqual(type(pattern), np.ndarray) |
| self.assertEqual(pattern.ndim, 2) |
| self.assertEqual(pattern.shape[0], pattern.shape[1]) |
| self.assertEqual(pattern.shape[0], workspace.NumGpuDevices()) |
| |
| @unittest.skipIf( |
| not workspace.has_cuda_support, "Tensor interop doesn't yet work on ROCm" |
| ) |
| def testTorchInterop(self): |
| # CUDA has convenient mem stats, let's use them to make sure we didn't |
| # leak memory |
| initial_mem = torch.cuda.memory_allocated() |
| workspace.RunOperatorOnce( |
| core.CreateOperator( |
| "ConstantFill", |
| [], |
| "foo", |
| shape=(4,), |
| value=2, |
| dtype=10, |
| device_option=core.DeviceOption(workspace.GpuDeviceType), |
| ) |
| ) |
| t = workspace.FetchTorch("foo") |
| t.resize_(5) |
| self.assertTrue(t.is_cuda) |
| t[4] = t[2] = 777 |
| np.testing.assert_array_equal(t.cpu().numpy(), np.array([2, 2, 777, 2, 777])) |
| np.testing.assert_array_equal( |
| workspace.FetchBlob("foo"), np.array([2, 2, 777, 2, 777]) |
| ) |
| |
| z = torch.ones((4,), dtype=torch.int64, device="cuda") |
| workspace.FeedBlob("bar", z) |
| workspace.RunOperatorOnce( |
| core.CreateOperator( |
| "Reshape", |
| ["bar"], |
| ["bar", "_"], |
| shape=(2, 2), |
| device_option=core.DeviceOption(workspace.GpuDeviceType), |
| ) |
| ) |
| z[0, 1] = 123 |
| np.testing.assert_array_equal( |
| workspace.FetchBlob("bar"), np.array([[1, 123], [1, 1]]) |
| ) |
| np.testing.assert_array_equal(z.cpu(), np.array([[1, 123], [1, 1]])) |
| |
| self.assertGreater(torch.cuda.memory_allocated(), initial_mem) |
| # clean up everything |
| del t |
| del z |
| workspace.ResetWorkspace() |
| self.assertEqual(torch.cuda.memory_allocated(), initial_mem) |
| |
| |
| @unittest.skipIf(not workspace.C.use_mkldnn, "No MKLDNN support.") |
| class TestWorkspaceIDEEP(test_util.TestCase): |
| def testFeedFetchBlobIDEEP(self): |
| arr = np.random.randn(2, 3).astype(np.float32) |
| workspace.FeedBlob("testblob_ideep", arr, core.DeviceOption(caffe2_pb2.IDEEP)) |
| fetched = workspace.FetchBlob("testblob_ideep") |
| np.testing.assert_array_equal(arr, fetched) |
| |
| |
| class TestImmedibate(test_util.TestCase): |
| def testImmediateEnterExit(self): |
| workspace.StartImmediate(i_know=True) |
| self.assertTrue(workspace.IsImmediate()) |
| workspace.StopImmediate() |
| self.assertFalse(workspace.IsImmediate()) |
| |
| def testImmediateRunsCorrectly(self): |
| workspace.StartImmediate(i_know=True) |
| net = core.Net("test-net") |
| net.ConstantFill([], "testblob", shape=[1, 2, 3, 4], value=1.0) |
| self.assertEqual(workspace.ImmediateBlobs(), ["testblob"]) |
| content = workspace.FetchImmediate("testblob") |
| # Also, the immediate mode should not invade the original namespace, |
| # so we check if this is so. |
| with self.assertRaises(RuntimeError): |
| workspace.FetchBlob("testblob") |
| np.testing.assert_array_equal(content, 1.0) |
| content[:] = 2.0 |
| self.assertTrue(workspace.FeedImmediate("testblob", content)) |
| np.testing.assert_array_equal(workspace.FetchImmediate("testblob"), 2.0) |
| workspace.StopImmediate() |
| with self.assertRaises(RuntimeError): |
| content = workspace.FetchImmediate("testblob") |
| |
| def testImmediateRootFolder(self): |
| workspace.StartImmediate(i_know=True) |
| # for testing we will look into the _immediate_root_folder variable |
| # but in normal usage you should not access that. |
| self.assertTrue(len(workspace._immediate_root_folder) > 0) |
| root_folder = workspace._immediate_root_folder |
| self.assertTrue(os.path.isdir(root_folder)) |
| workspace.StopImmediate() |
| self.assertTrue(len(workspace._immediate_root_folder) == 0) |
| # After termination, immediate mode should have the root folder |
| # deleted. |
| self.assertFalse(os.path.exists(root_folder)) |
| |
| |
| class TestCppEnforceAsException(test_util.TestCase): |
| def testEnforce(self): |
| op = core.CreateOperator("Relu", ["X"], ["Y"]) |
| with self.assertRaises(RuntimeError): |
| workspace.RunOperatorOnce(op) |
| |
| |
| class TestCWorkspace(htu.HypothesisTestCase): |
| def test_net_execution(self): |
| ws = workspace.C.Workspace() |
| self.assertEqual(ws.nets, {}) |
| self.assertEqual(ws.blobs, {}) |
| net = core.Net("test-net") |
| net.ConstantFill([], "testblob", shape=[1, 2, 3, 4], value=1.0) |
| ws.create_net(net) |
| # If we do not specify overwrite, this should raise an error. |
| with self.assertRaises(RuntimeError): |
| ws.create_net(net) |
| # But, if we specify overwrite, this should pass. |
| ws.create_net(net, True) |
| # Overwrite can also be a kwarg. |
| ws.create_net(net, overwrite=True) |
| self.assertIn("testblob", ws.blobs) |
| self.assertEqual(len(ws.nets), 1) |
| net_name = net.Proto().name |
| self.assertIn("test-net", net_name) |
| net = ws.nets[net_name].run() |
| blob = ws.blobs["testblob"] |
| np.testing.assert_array_equal( |
| np.ones((1, 2, 3, 4), dtype=np.float32), blob.fetch() |
| ) |
| |
| @given(name=st.text(), value=st.floats(min_value=-1, max_value=1.0)) |
| def test_operator_run(self, name, value): |
| ws = workspace.C.Workspace() |
| op = core.CreateOperator("ConstantFill", [], [name], shape=[1], value=value) |
| ws.run(op) |
| self.assertIn(name, ws.blobs) |
| np.testing.assert_allclose( |
| [value], ws.blobs[name].fetch(), atol=1e-4, rtol=1e-4 |
| ) |
| |
| @given( |
| blob_name=st.text(), |
| net_name=st.text(), |
| value=st.floats(min_value=-1, max_value=1.0), |
| ) |
| def test_net_run(self, blob_name, net_name, value): |
| ws = workspace.C.Workspace() |
| net = core.Net(net_name) |
| net.ConstantFill([], [blob_name], shape=[1], value=value) |
| ws.run(net) |
| self.assertIn(blob_name, ws.blobs) |
| self.assertNotIn(net_name, ws.nets) |
| np.testing.assert_allclose( |
| [value], ws.blobs[blob_name].fetch(), atol=1e-4, rtol=1e-4 |
| ) |
| |
| @given( |
| blob_name=st.text(), |
| net_name=st.text(), |
| plan_name=st.text(), |
| value=st.floats(min_value=-1, max_value=1.0), |
| ) |
| def test_plan_run(self, blob_name, plan_name, net_name, value): |
| ws = workspace.C.Workspace() |
| plan = core.Plan(plan_name) |
| net = core.Net(net_name) |
| net.ConstantFill([], [blob_name], shape=[1], value=value) |
| |
| plan.AddStep(core.ExecutionStep("step", nets=[net], num_iter=1)) |
| |
| ws.run(plan) |
| self.assertIn(blob_name, ws.blobs) |
| self.assertIn(net.Name(), ws.nets) |
| np.testing.assert_allclose( |
| [value], ws.blobs[blob_name].fetch(), atol=1e-4, rtol=1e-4 |
| ) |
| |
| @given( |
| blob_name=st.text(), |
| net_name=st.text(), |
| value=st.floats(min_value=-1, max_value=1.0), |
| ) |
| def test_net_create(self, blob_name, net_name, value): |
| ws = workspace.C.Workspace() |
| net = core.Net(net_name) |
| net.ConstantFill([], [blob_name], shape=[1], value=value) |
| ws.create_net(net).run() |
| self.assertIn(blob_name, ws.blobs) |
| self.assertIn(net.Name(), ws.nets) |
| np.testing.assert_allclose( |
| [value], ws.blobs[blob_name].fetch(), atol=1e-4, rtol=1e-4 |
| ) |
| |
| @given( |
| name=st.text(), |
| value=htu.tensor(), |
| device_option=st.sampled_from(htu.device_options), |
| ) |
| def test_array_serde(self, name, value, device_option): |
| ws = workspace.C.Workspace() |
| ws.create_blob(name).feed(value, device_option=device_option) |
| self.assertIn(name, ws.blobs) |
| blob = ws.blobs[name] |
| np.testing.assert_equal(value, ws.blobs[name].fetch()) |
| serde_blob = ws.create_blob("{}_serde".format(name)) |
| serde_blob.deserialize(blob.serialize(name)) |
| np.testing.assert_equal(value, serde_blob.fetch()) |
| |
| @given(name=st.text(), value=st.text()) |
| def test_string_serde(self, name, value): |
| value = value.encode("ascii", "ignore") |
| ws = workspace.C.Workspace() |
| ws.create_blob(name).feed(value) |
| self.assertIn(name, ws.blobs) |
| blob = ws.blobs[name] |
| self.assertEqual(value, ws.blobs[name].fetch()) |
| serde_blob = ws.create_blob("{}_serde".format(name)) |
| serde_blob.deserialize(blob.serialize(name)) |
| self.assertEqual(value, serde_blob.fetch()) |
| |
| def test_exception(self): |
| ws = workspace.C.Workspace() |
| |
| with self.assertRaises(TypeError): |
| ws.create_net("...") |
| |
| |
| class TestPredictor(unittest.TestCase): |
| def _create_model(self): |
| m = model_helper.ModelHelper() |
| y = brew.fc( |
| m, |
| "data", |
| "y", |
| dim_in=4, |
| dim_out=2, |
| weight_init=("ConstantFill", dict(value=1.0)), |
| bias_init=("ConstantFill", dict(value=0.0)), |
| axis=0, |
| ) |
| m.net.AddExternalOutput(y) |
| return m |
| |
| # Use this test with a bigger model to see how using Predictor allows to |
| # avoid issues with low protobuf size limit in Python |
| # |
| # def test_predictor_predefined(self): |
| # workspace.ResetWorkspace() |
| # path = 'caffe2/caffe2/test/assets/' |
| # with open(path + 'squeeze_predict_net.pb') as f: |
| # self.predict_net = f.read() |
| # with open(path + 'squeeze_init_net.pb') as f: |
| # self.init_net = f.read() |
| # self.predictor = workspace.Predictor(self.init_net, self.predict_net) |
| |
| # inputs = [np.zeros((1, 3, 256, 256), dtype='f')] |
| # outputs = self.predictor.run(inputs) |
| # self.assertEqual(len(outputs), 1) |
| # self.assertEqual(outputs[0].shape, (1, 1000, 1, 1)) |
| # self.assertAlmostEqual(outputs[0][0][0][0][0], 5.19026289e-05) |
| |
| def test_predictor_memory_model(self): |
| workspace.ResetWorkspace() |
| m = self._create_model() |
| workspace.FeedBlob("data", np.zeros([4], dtype="float32")) |
| self.predictor = workspace.Predictor( |
| workspace.StringifyProto(m.param_init_net.Proto()), |
| workspace.StringifyProto(m.net.Proto()), |
| ) |
| |
| inputs = np.array([1, 3, 256, 256], dtype="float32") |
| outputs = self.predictor.run([inputs]) |
| np.testing.assert_array_almost_equal( |
| np.array([[516, 516]], dtype="float32"), outputs |
| ) |
| |
| |
| class TestTransform(htu.HypothesisTestCase): |
| @given( |
| input_dim=st.integers(min_value=1, max_value=10), |
| output_dim=st.integers(min_value=1, max_value=10), |
| batch_size=st.integers(min_value=1, max_value=10), |
| ) |
| def test_simple_transform(self, input_dim, output_dim, batch_size): |
| m = model_helper.ModelHelper() |
| fc1 = brew.fc(m, "data", "fc1", dim_in=input_dim, dim_out=output_dim) |
| fc2 = brew.fc(m, fc1, "fc2", dim_in=output_dim, dim_out=output_dim) |
| conv = brew.conv( |
| m, |
| fc2, |
| "conv", |
| dim_in=output_dim, |
| dim_out=output_dim, |
| use_cudnn=True, |
| engine="CUDNN", |
| kernel=3, |
| ) |
| |
| conv.Relu([], conv).Softmax([], "pred").LabelCrossEntropy( |
| ["label"], ["xent"] |
| ).AveragedLoss([], "loss") |
| |
| transformed_net_proto = workspace.ApplyTransform("ConvToNNPack", m.net.Proto()) |
| |
| self.assertEqual(transformed_net_proto.op[2].engine, "NNPACK") |
| |
| @given( |
| input_dim=st.integers(min_value=1, max_value=10), |
| output_dim=st.integers(min_value=1, max_value=10), |
| batch_size=st.integers(min_value=1, max_value=10), |
| ) |
| @settings(deadline=10000) |
| def test_registry_invalid(self, input_dim, output_dim, batch_size): |
| m = model_helper.ModelHelper() |
| brew.fc(m, "data", "fc1", dim_in=input_dim, dim_out=output_dim) |
| with self.assertRaises(RuntimeError): |
| workspace.ApplyTransform("definitely_not_a_real_transform", m.net.Proto()) |
| |
| @given(value=st.floats(min_value=-1, max_value=1)) |
| @settings(deadline=10000) |
| def test_apply_transform_if_faster(self, value): |
| |
| init_net = core.Net("init_net") |
| init_net.ConstantFill([], ["data"], shape=[5, 5, 5, 5], value=value) |
| init_net.ConstantFill([], ["conv_w"], shape=[5, 5, 3, 3], value=value) |
| init_net.ConstantFill([], ["conv_b"], shape=[5], value=value) |
| |
| self.assertEqual( |
| workspace.RunNetOnce(init_net.Proto().SerializeToString()), True |
| ) |
| |
| m = model_helper.ModelHelper() |
| conv = brew.conv( |
| m, |
| "data", |
| "conv", |
| dim_in=5, |
| dim_out=5, |
| kernel=3, |
| use_cudnn=True, |
| engine="CUDNN", |
| ) |
| |
| conv.Relu([], conv).Softmax([], "pred").AveragedLoss([], "loss") |
| |
| self.assertEqual(workspace.RunNetOnce(m.net.Proto().SerializeToString()), True) |
| |
| proto = workspace.ApplyTransformIfFaster( |
| "ConvToNNPack", m.net.Proto(), init_net.Proto() |
| ) |
| self.assertEqual(workspace.RunNetOnce(proto.SerializeToString()), True) |
| proto = workspace.ApplyTransformIfFaster( |
| "ConvToNNPack", |
| m.net.Proto(), |
| init_net.Proto(), |
| warmup_runs=10, |
| main_runs=100, |
| improvement_threshold=2.0, |
| ) |
| self.assertEqual(workspace.RunNetOnce(proto.SerializeToString()), True) |
| |
| |
| class MyModule(torch.jit.ScriptModule): |
| def __init__(self): |
| super(MyModule, self).__init__() |
| self.mult = torch.nn.Parameter(torch.tensor([[1, 2, 3, 4, 5.0]])) |
| |
| @torch.jit.script_method |
| def forward(self, x): |
| return self.mult.mm(x) |
| |
| @torch.jit.script_method |
| def multi_input(self, x: torch.Tensor, y: torch.Tensor, z: int = 2) -> torch.Tensor: |
| return x + y + z |
| |
| @torch.jit.script_method |
| def multi_input_tensor_list(self, tensor_list: List[Tensor]) -> Tensor: |
| return tensor_list[0] + tensor_list[1] + tensor_list[2] |
| |
| @torch.jit.script_method |
| def multi_output(self, x): |
| return (x, x + 1) |
| |
| |
| @unittest.skipIf( |
| "ScriptModule" not in core._REGISTERED_OPERATORS, |
| "Script module integration in Caffe2 is not enabled", |
| ) |
| class TestScriptModule(test_util.TestCase): |
| def _createFeedModule(self): |
| workspace.FeedBlob("m", MyModule()) |
| |
| def testCreation(self): |
| m = MyModule() |
| workspace.FeedBlob("module", m) |
| m2 = workspace.FetchBlob("module") |
| self.assertTrue(m2 is not None) |
| |
| def testForward(self): |
| self._createFeedModule() |
| val = np.random.rand(5, 5).astype(np.float32) |
| param = np.array([[1, 2, 3, 4, 5]]).astype(np.float32) |
| workspace.FeedBlob("w", val) |
| workspace.RunOperatorOnce( |
| core.CreateOperator("ScriptModule", ["m", "w"], ["y"]) |
| ) |
| np.testing.assert_almost_equal( |
| workspace.FetchBlob("y"), np.matmul(param, val), decimal=5 |
| ) |
| |
| def testMultiInputOutput(self): |
| self._createFeedModule() |
| val = np.random.rand(5, 5).astype(np.float32) |
| workspace.FeedBlob("w", val) |
| val2 = np.random.rand(5, 5).astype(np.float32) |
| workspace.FeedBlob("w2", val2) |
| workspace.RunOperatorOnce( |
| core.CreateOperator( |
| "ScriptModule", ["m", "w", "w2"], ["y"], method="multi_input" |
| ) |
| ) |
| workspace.RunOperatorOnce( |
| core.CreateOperator( |
| "ScriptModule", ["m", "w"], ["y1", "y2"], method="multi_output" |
| ) |
| ) |
| np.testing.assert_almost_equal( |
| workspace.FetchBlob("y"), val + val2 + 2, decimal=5 |
| ) |
| np.testing.assert_almost_equal(workspace.FetchBlob("y1"), val, decimal=5) |
| np.testing.assert_almost_equal(workspace.FetchBlob("y2"), val + 1, decimal=5) |
| |
| def testMultiTensorListInput(self): |
| self._createFeedModule() |
| val = np.random.rand(5, 5).astype(np.float32) |
| workspace.FeedBlob("w", val) |
| val2 = np.random.rand(5, 5).astype(np.float32) |
| workspace.FeedBlob("w2", val2) |
| val3 = np.random.rand(5, 5).astype(np.float32) |
| workspace.FeedBlob("w3", val3) |
| |
| workspace.RunOperatorOnce( |
| core.CreateOperator( |
| "ScriptModule", |
| ["m", "w", "w2", "w3"], |
| ["y"], |
| method="multi_input_tensor_list", |
| pass_inputs_as_tensor_list=True, |
| ) |
| ) |
| np.testing.assert_almost_equal( |
| workspace.FetchBlob("y"), val + val2 + val3, decimal=5 |
| ) |
| |
| def testSerialization(self): |
| tmpdir = tempfile.mkdtemp() |
| try: |
| self._createFeedModule() |
| workspace.RunOperatorOnce( |
| core.CreateOperator( |
| "Save", |
| ["m"], |
| [], |
| absolute_path=1, |
| db=os.path.join(tmpdir, "db"), |
| db_type="minidb", |
| ) |
| ) |
| workspace.ResetWorkspace() |
| |
| self.assertFalse(workspace.HasBlob("m")) |
| workspace.RunOperatorOnce( |
| core.CreateOperator( |
| "Load", |
| [], |
| [], |
| absolute_path=1, |
| db=os.path.join(tmpdir, "db"), |
| db_type="minidb", |
| load_all=1, |
| ) |
| ) |
| self.assertTrue(workspace.HasBlob("m")) |
| # TODO: make caffe2 side load return python-sided module |
| # right now it returns the base class (torch._C.ScriptModule) |
| # self.assertTrue(isinstance(workspace.FetchBlob('m'), torch.jit.ScriptModule)) |
| |
| # do something with the module |
| val = np.random.rand(5, 5).astype(np.float32) |
| param = np.array([[1, 2, 3, 4, 5]]).astype(np.float32) |
| workspace.FeedBlob("w", val) |
| workspace.RunOperatorOnce( |
| core.CreateOperator("ScriptModule", ["m", "w"], ["y"]) |
| ) |
| np.testing.assert_almost_equal( |
| workspace.FetchBlob("y"), np.matmul(param, val), decimal=5 |
| ) |
| finally: |
| # clean up temp folder. |
| try: |
| shutil.rmtree(tmpdir) |
| except OSError as e: |
| if e.errno != errno.ENOENT: |
| raise |
| |
| |
| class TestScriptModuleFromString(TestScriptModule): |
| def _createFeedModule(self): |
| workspace.RunOperatorOnce( |
| core.CreateOperator( |
| "ScriptModuleLoad", |
| [], |
| ["m"], |
| serialized_binary=self._get_modules_bytes(MyModule()), |
| ) |
| ) |
| |
| def _get_modules_bytes(self, the_module): |
| import io |
| |
| buffer = io.BytesIO() |
| torch.jit.save(the_module, buffer) |
| return buffer.getvalue() |
| |
| |
| if __name__ == "__main__": |
| unittest.main() |