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)