def verify_split(input_dim, nOutputs):
dtype = "float32"
a_np = np.random.uniform(-100.0, 100.0, size=input_dim).astype(dtype)
ref_val = np.split(a_np, nOutputs, axis=-3)
inputs = [("input", datatypes.Array(*input_dim))]
output_names = []
outputs = []
output_shapes = []
for i, out in enumerate(ref_val):
output_name = "output" + str(i)
output_names = output_names + [output_name]
outputs = outputs + [(output_name, datatypes.Array(*out.shape))]
output_shapes = output_shapes + [out.shape]
builder = NeuralNetworkBuilder(inputs, outputs)
builder.add_split(name="split", input_name="input", output_names=output_names)
model = cm.models.MLModel(builder.spec)
for target, dev in tvm.testing.enabled_targets():
out = run_tvm_graph(
model, target, dev, [a_np], ["input"], output_shapes, [dtype] * len(output_shapes)
)
tvm.testing.assert_allclose(out, ref_val, rtol=1e-5)
def test_split_converter(self):
input_dim = (8, 1, 1)
output_dim = (4, 1, 1)
inputs = [('input', datatypes.Array(*input_dim))]
outputs = [('output1', datatypes.Array(*output_dim)), ('output2', datatypes.Array(*output_dim))]
builder = NeuralNetworkBuilder(inputs, outputs)
builder.add_split(name='Split', input_name='input', output_names=['output1', 'output2'])
model_onnx = convert_coreml(builder.spec)
self.assertTrue(model_onnx is not None)
def test_split_converter(self):
input_dim = (8, 1, 1)
output_dim = (4, 1, 1)
inputs = [('input', datatypes.Array(*input_dim))]
outputs = [('output1', datatypes.Array(*output_dim)),
('output2', datatypes.Array(*output_dim))]
builder = NeuralNetworkBuilder(inputs, outputs)
builder.add_split(name='Split',
input_name='input',
output_names=['output1', 'output2'])
context = ConvertContext()
node = SplitLayerConverter.convert(
context, builder.spec.neuralNetwork.layers[0], ['input'],
['output'])
self.assertTrue(node is not None)
