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)