def test_multiple_registrations(self):
assert OpaqueFuncRegistry.Size() == 0
@register_opaque_func('py_func1', [TypeCode.Str, TypeCode.XGraph])
def py_func(str_name, xgraph):
assert xgraph.get_name() == "name1"
xgraph.set_name(str_name)
@register_opaque_func('py_func2', [TypeCode.Str, TypeCode.XGraph])
def py_func2(str_name, xgraph):
xgraph.add(XLayer(name=str_name, type=[str_name]))
assert OpaqueFuncRegistry.Size() == 2
assert set(OpaqueFuncRegistry.GetRegisteredFuncs()) == \
set(['py_func1', 'py_func2'])
xg = XGraph("name1")
of1 = OpaqueFuncRegistry.Get("py_func1")
of1("name2", xg)
assert xg.get_name() == "name2"
assert len(xg) == 0
of2 = OpaqueFuncRegistry.Get("py_func2")
of2("x1", xg)
assert xg.get_name() == "name2"
assert len(xg) == 1
assert xg.get('x1').name == 'x1'
assert xg.get('x1').type == ['x1']
def test_registration_flow(self):
assert OpaqueFuncRegistry.Size() == 0
@register_opaque_func('py_func', [TypeCode.Str, TypeCode.XGraph])
def py_func(str_name, xgraph):
assert xgraph.get_name() == "name1"
xgraph.set_name(str_name)
assert OpaqueFuncRegistry.Size() == 1
assert OpaqueFuncRegistry.GetRegisteredFuncs() == ['py_func']
xg = XGraph("name1")
of = OpaqueFuncRegistry.Get("py_func")
of("name2", xg)
assert xg.get_name() == "name2"
def test_basic_flow(self):
def py_func(str_name):
assert str_name == "test"
of = OpaqueFunc(py_func, [TypeCode.Str])
assert OpaqueFuncRegistry.Size() == 0
ofr = OpaqueFuncRegistry.Register('py_func')
ofr.set_func(of)
assert OpaqueFuncRegistry.Size() == 1
assert OpaqueFuncRegistry.GetRegisteredFuncs() == ['py_func']
of2 = OpaqueFuncRegistry.Get("py_func")
of2("test")
def read_xgraph_str(xg_str: bytes):
of = OpaqueFuncRegistry.Get("pyxir.io.deserialize_xgraph")
xg = XGraph()
s = BytesContainer(xg_str)
# import pdb; pdb.set_trace()
of(xg, s)
return xg
def get_xgraph_str(xg: XGraph):
# graph_str, data_str = XGraphIO.to_string(xg)
# return " " + str(len(graph_str)) + " " + graph_str + " " + str(len(data_str) + 1) + " " + data_str
of = OpaqueFuncRegistry.Get("pyxir.io.get_serialized_xgraph")
s = BytesContainer(b"")
of(xg, s)
# import pdb; pdb.set_trace()
return s.get_bytes()
def test_empty_directory_serialization(self):
dir_path = "/tmp/test_dir"
if os.path.exists(dir_path):
shutil.rmtree(dir_path)
os.makedirs(dir_path)
bytes_c = BytesContainer(b"")
of = OpaqueFuncRegistry.Get("pyxir.io.serialize_dir")
of(dir_path, bytes_c)
shutil.rmtree("/tmp/test_dir")
assert not os.path.exists("/tmp/test_dir")
of_de = OpaqueFuncRegistry.Get("pyxir.io.deserialize_dir")
of_de(dir_path, bytes_c.get_bytes())
assert os.path.exists("/tmp/test_dir")
def test_directory_serialization(self):
dir_path = "/tmp/test_dir"
if os.path.exists(dir_path):
shutil.rmtree(dir_path)
os.makedirs(dir_path)
with open("/tmp/test_dir/test.txt", 'w') as f:
f.write("testtest")
bytes_c = BytesContainer(b"")
of = OpaqueFuncRegistry.Get("pyxir.io.serialize_dir")
of(dir_path, bytes_c)
shutil.rmtree("/tmp/test_dir")
assert not os.path.exists("/tmp/test_dir")
of_de = OpaqueFuncRegistry.Get("pyxir.io.deserialize_dir")
of_de(dir_path, bytes_c.get_bytes())
assert os.path.exists("/tmp/test_dir")
assert os.path.exists("/tmp/test_dir/test.txt")
with open("/tmp/test_dir/test.txt", 'r') as f:
assert f.read() == "testtest"
def test_registration_flow_same_name(self):
assert OpaqueFuncRegistry.Size() == 0
@register_opaque_func('py_func', [TypeCode.Str, TypeCode.XGraph])
def py_func(str_name, xgraph):
assert xgraph.get_name() == "name1"
xgraph.set_name(str_name)
with self.assertRaises(ValueError):
@register_opaque_func('py_func', [TypeCode.Str, TypeCode.XGraph])
def py_func(str_name, xgraph):
assert xgraph.get_name() == "name1"
xgraph.set_name(str_name)
def test_build_rt_opaque_func_cpu_tf(self):
tf.compat.v1.reset_default_graph()
W = np.reshape(
np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], dtype=np.float32),
(2, 1, 2, 2))
# B = np.array([0, 0], dtype=np.float32)
iX = px.ops.input('in', [1, 1, 4, 4])
wX = px.ops.constant('w', W)
# bX = px.ops.constant('b', B)
cX = px.ops.conv2d('conv',
iX,
wX,
kernel_size=[2, 2],
strides=[1, 1],
padding_hw=[0, 0, 0, 0],
dilation=[1, 1],
groups=1,
channels=2,
data_layout='NCHW',
kernel_layout='OIHW')
xlayers = [iX, wX, cX]
xg = TestBase.xg_factory.build_from_xlayer(xlayers)
of = OpaqueFuncRegistry.Get('pyxir.build_rt')
rt_of = OpaqueFunc()
# call opaque func
of(xg, 'cpu', 'cpu-tf', ['in'], ['conv'], rt_of)
# By now, the `rt_of` is initialized with the opaque runtime function
ins = [XBuffer(np.ones((1, 1, 4, 4), dtype=np.float32))]
outs = [XBuffer(np.empty((1, 2, 3, 3), dtype=np.float32))]
rt_of(ins, outs)
assert len(outs) == 1
expected_outpt = np.array([[[[10., 10., 10.], [10., 10., 10.],
[10., 10., 10.]],
[[26., 26., 26.], [26., 26., 26.],
[26., 26., 26.]]]])
np.testing.assert_array_equal(outs[0], expected_outpt)
def test_simple_model_opaque_func(self):
x = helper.make_tensor_value_info('x', TensorProto.FLOAT,
[None, 1, 4, 4])
x_val = np.array([[[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12],
[13, 14, 15, 16]]]]).astype(np.float32)
# x_init = helper.make_tensor('x', TensorProto.FLOAT, (1, 1, 4, 4),
# list(x_val.reshape(-1)))
# Create one output (ValueInfoProto)
z = helper.make_tensor_value_info('z', TensorProto.FLOAT,
[None, 2, 2, 2])
W_val = np.array([[[[1, 1], [1, 1]]], [[[1, -1],
[1, 1]]]]).astype(np.float32)
W = helper.make_tensor('W', TensorProto.FLOAT, (2, 1, 2, 2),
list(W_val.reshape(-1)))
B_val = np.array([1, -1]).astype(np.float32)
B = helper.make_tensor('B', TensorProto.FLOAT, (2, ),
list(B_val.reshape((-1))))
conv_node = onnx.helper.make_node('Conv',
inputs=['x', 'W', 'B'],
outputs=['y'],
kernel_shape=[2, 2],
pads=[1, 1, 0, 0])
pool_node = onnx.helper.make_node('AveragePool',
inputs=['y'],
outputs=['z'],
kernel_shape=[2, 2],
pads=[0, 0, 0, 0],
strides=[2, 2])
# Create the graph (GraphProto)
graph_def = onnx.helper.make_graph(
[conv_node, pool_node],
'test-model',
[x],
[z],
[W, B] # x_init]
)
# Create the model (ModelProto)
model_def = onnx.helper.make_model(graph_def,
producer_name='onnx-example')
test_file = os.path.join(FILE_DIR, 'test.onnx')
onnx.save(model_def, test_file)
xgraph = XGraph(name='test')
of = OpaqueFuncRegistry.Get('pyxir.onnx.from_onnx')
of(xgraph, test_file)
assert xgraph.get_name() == 'test-model'
xlayers = xgraph.get_layers()
assert len(xlayers) == 4
assert xlayers[0].name == 'x'
assert xlayers[0].type[0] == 'Input'
assert xlayers[0].shapes == [-1, 1, 4, 4]
assert xlayers[0].attrs['onnx_id'] == 'x'
assert xlayers[1].name == 'y_Conv'
assert xlayers[1].type[0] == 'Convolution'
assert xlayers[1].shapes == [-1, 2, 4, 4]
assert xlayers[1].attrs['padding'] == [(0, 0), (0, 0), (1, 0), (1, 0)]
assert xlayers[1].attrs['strides'] == [1, 1]
assert xlayers[1].attrs['dilation'] == [1, 1]
assert xlayers[1].attrs['kernel_size'] == [2, 2]
assert xlayers[1].attrs['channels'] == [1, 2]
assert xlayers[1].attrs['data_layout'] == 'NCHW'
assert xlayers[1].attrs['kernel_layout'] == 'OIHW'
assert xlayers[1].attrs['groups'] == 1
assert xlayers[1].attrs['onnx_id'] == 'y'
assert xlayers[2].name == 'y'
assert xlayers[2].shapes == [-1, 2, 4, 4]
assert xlayers[2].attrs['axis'] == 1
assert xlayers[2].attrs['onnx_id'] == 'y'
assert xlayers[3].name == 'z'
assert xlayers[3].shapes == [-1, 2, 2, 2]
assert xlayers[3].type[0] == 'Pooling'
assert xlayers[3].attrs['padding'] == [[0, 0], [0, 0], [0, 0], [0, 0]]
assert xlayers[3].attrs['strides'] == [2, 2]
assert xlayers[3].attrs['kernel_size'] == [2, 2]
assert xlayers[3].attrs['data_layout'] == 'NCHW'
assert xlayers[3].attrs['type'] == 'Avg'
assert xlayers[3].attrs['onnx_id'] == 'z'
of = OpaqueFuncRegistry.Get('pyxir.partition')
of(xgraph, ['test_dpu'], "")
assert xgraph.get_name() == 'test-model'
assert len(xgraph) == 4
xlayers = xgraph.get_layers()
assert xlayers[0].name == 'x'
assert xlayers[0].target == 'cpu'
assert xlayers[0].subgraph is None
assert xlayers[1].name == 'y_Conv'
assert xlayers[1].target == 'test_dpu'
assert xlayers[1].subgraph == 'xp0'
assert xlayers[2].name == 'y'
assert xlayers[2].type == ['BiasAdd']
assert xlayers[2].target == 'test_dpu'
assert xlayers[2].subgraph == 'xp0'
assert xlayers[3].name == 'z'
assert xlayers[3].type == ['Pooling']
assert xlayers[3].target == 'test_dpu'
assert xlayers[3].subgraph == 'xp0'
os.remove(test_file)
def is_dpuczdx8g_vart_flow_enabled():
of = OpaqueFuncRegistry.Get("pyxir.use_dpuczdx8g_vart")
s = StrContainer("")
of(s)
return s.get_str() == "True"
def setUp(self):
OpaqueFuncRegistry.Clear()