pytorch/caffe2/python/workspace_test.py
Edward Yang 25261a4776 Merge Tensor and Variable. (#28620)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/28620

All Tensors are Variables now, they just happen to have requires_grad=False. Tensors ALWAYS have `VariableTensorId` in their type set.

When constructing this patch, I had to make decisions about what I would fix in this patch, and what I would leave for follow up PRs. Here is the cleanup that happens in this patch:

- The `is_variable` property is removed from TensorOptions. I removed this immediately because unlike Tensor::is_variable, TensorOptions::is_variable doesn't respect our VariableTensorId thread-local state. This means that there were a bunch of places where TensorOptions::is_variable was false, which is obviously bogus in the world when tensor and variable are merged. Instead of keeping the method as a function that always returns true, I just opted to remove it entirely (it's not public API.) All places we set `is_variable` are deleted.
  - Knock on effect: there is no longer a separate DeprecatedTypeProperties for the variable and non-variable versions of type.
  - Knock on effect: instead of asserting on TensorOptions::is_variable, instead we just test `at::impl::variable_is_excluded()`
- There is now only one copy of the cuDNN RNN dropout cache, not two (I'm not sure why we had two to begin with)

Some cleanup that doesn't happen in this patch:
- Eliminating unnecessary uses of `make_variable`
- Eliminating `Tensor::is_variable`

The most subtle part of this patch is retaining tracing behavior: the fact that everything is a Variable means that more code gets routed to VariableType than before; this can change traces. I identified two places where we didn't appropriately turn off VariableType, mostly factory functions:

- `torch.tensor` must turn off VariableType before invoking `at::empty` to construct the tensor, as it subsequently does direct data access
- `tensor_slow` (invoked when you pass a Python scalar to a tensor argument) must turn off VariableType before calling `scalar_to_tensor` so the scalar gets traced as constant, rather than as a call to `scalar_to_tensor`.

Honestly, these are all giant hacks, and should be replaced with a more specialized guard that just toggles tracing.

Signed-off-by: Edward Z. Yang <ezyang@fb.com>

Test Plan: Imported from OSS

Reviewed By: dreiss

Differential Revision: D18171156

Pulled By: ezyang

fbshipit-source-id: 5b6a045beba37492647e350190f495114e86504d
2019-11-04 14:59:57 -08:00

808 lines
32 KiB
Python

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
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
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))
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))
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, y, z=2):
# type: (Tensor, Tensor, int) -> 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()