frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 13:44:15 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
dba94873246ba71c40fbe98e6a06e23d6d59ee86
pytorch/caffe2/python/workspace_test.py
Xuehai Pan 8d45f555d7 [BE] [1/3] Rewrite super() calls in caffe2 and benchmarks (#94587) 
Rewrite Python built-in class `super()` calls. Only non-semantic changes should be applied.

- #94587
- #94588
- #94592

Also, methods with only a `super()` call are removed:

```diff
class MyModule(nn.Module):
-   def __init__(self):
-       super().__init__()
-
    def forward(self, ...):
        ...
```

Some cases that change the semantics should be kept unchanged. E.g.:

f152a79be9/caffe2/python/net_printer.py (L184-L190)

f152a79be9/test/test_jit_fuser_te.py (L2628-L2635)

Pull Request resolved: https://github.com/pytorch/pytorch/pull/94587
Approved by: https://github.com/ezyang
2023-02-11 18:19:48 +00:00

934 lines
34 KiB
Python
Raw Blame History

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.assertEqual(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.assertEqual(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.assertEqual(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.assertEqual(s1, fetch1)
self.assertEqual(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().__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()
Reference in New Issue View Git Blame Copy Permalink