mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/52083 This makes minor fixes in `caffe2/python` to address all errors currently reported by Pyre. I update the code to fix errors when doing so looked simple and safe, and added `pyre-fixme` comments in other places. ghstack-source-id: 121109695 Test Plan: Confirmed that Pyre no longer reports errors under `caffe2/python` Differential Revision: D26272279 fbshipit-source-id: b1eb19d323b613f23280ce9c71e800e874ca1162
811 lines
32 KiB
Python
811 lines
32 KiB
Python
|
|
|
|
|
|
|
|
|
|
import numpy as np
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
import unittest
|
|
|
|
import torch
|
|
from caffe2.proto import caffe2_pb2
|
|
from caffe2.python import core, test_util, workspace, model_helper, brew
|
|
|
|
import caffe2.python.hypothesis_test_util as htu
|
|
import hypothesis.strategies as st
|
|
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 testRootFolder(self):
|
|
self.assertEqual(workspace.ResetWorkspace(), True)
|
|
self.assertEqual(workspace.RootFolder(), ".")
|
|
self.assertEqual(
|
|
workspace.ResetWorkspace("/tmp/caffe-workspace-test"), True)
|
|
self.assertEqual(workspace.RootFolder(), "/tmp/caffe-workspace-test")
|
|
|
|
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)
|
|
flops, _ = workspace.GetOperatorCost(op.SerializeToString(), ["X", "W"])
|
|
self.assertEqual(flops, 1152)
|
|
|
|
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)
|
|
|
|
# pyre-fixme[56]: Pyre was not able to infer the type of the decorator
|
|
# `torch.jit.script_method`
|
|
@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_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 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()
|