def test_simple_branch(self):
""" Test a simple if-else branch network
"""
input_features = [('data', datatypes.Array(3)), ('cond', datatypes.Array(1))]
output_features = [('output', None)]
builder_top = NeuralNetworkBuilder(input_features, output_features, disable_rank5_shape_mapping=True)
layer = builder_top.add_branch('branch_layer', 'cond')
builder_ifbranch = NeuralNetworkBuilder(input_features=None, output_features=None, spec=None, nn_spec=layer.branch.ifBranch)
builder_ifbranch.add_elementwise('mult_layer', input_names=['data'], output_name='output', mode='MULTIPLY', alpha=10)
builder_elsebranch = NeuralNetworkBuilder(input_features=None, output_features=None, spec=None, nn_spec=layer.branch.elseBranch)
builder_elsebranch.add_elementwise('add_layer', input_names=['data'], output_name='output', mode='ADD', alpha=10)
coremltools.models.utils.save_spec(builder_top.spec, '/tmp/simple_branch.mlmodel')
mlmodel = MLModel(builder_top.spec)
# True branch case
input_dict = {'data': np.array(range(1,4), dtype='float'), 'cond': np.array([1], dtype='float')}
output_ref = {'output': input_dict['data'] * 10}
self._test_model(mlmodel, input_dict, output_ref)
# False branch case
input_dict['cond'] = np.array([0], dtype='float')
output_ref['output'] = input_dict['data'] + 10
self._test_model(mlmodel, input_dict, output_ref)
def test_simple_branch(self):
""" Test a simple if-else branch network
"""
input_features = [("data", datatypes.Array(3)), ("cond", datatypes.Array(1))]
output_features = [("output", None)]
builder_top = NeuralNetworkBuilder(
input_features, output_features, disable_rank5_shape_mapping=True
)
layer = builder_top.add_branch("branch_layer", "cond")
builder_ifbranch = NeuralNetworkBuilder(
input_features=None,
output_features=None,
spec=None,
nn_spec=layer.branch.ifBranch,
)
builder_ifbranch.add_elementwise(
"mult_layer",
input_names=["data"],
output_name="output",
mode="MULTIPLY",
alpha=10,
)
builder_elsebranch = NeuralNetworkBuilder(
input_features=None,
output_features=None,
spec=None,
nn_spec=layer.branch.elseBranch,
)
builder_elsebranch.add_elementwise(
"add_layer",
input_names=["data"],
output_name="output",
mode="ADD",
alpha=10,
)
coremltools.models.utils.save_spec(
builder_top.spec, "/tmp/simple_branch.mlmodel"
)
mlmodel = MLModel(builder_top.spec)
# True branch case
input_dict = {
"data": np.array(range(1, 4), dtype="float"),
"cond": np.array([1], dtype="float"),
}
output_ref = {"output": input_dict["data"] * 10}
self._test_model(mlmodel, input_dict, output_ref)
# False branch case
input_dict["cond"] = np.array([0], dtype="float")
output_ref["output"] = input_dict["data"] + 10
self._test_model(mlmodel, input_dict, output_ref)
