def test_simple_loop_fixed_iterations(self):
input_features = [("data", datatypes.Array(1))]
output_features = [("output", None)]
builder_top = NeuralNetworkBuilder(
input_features, output_features, disable_rank5_shape_mapping=True
)
builder_top.add_copy("copy_1", input_name="data", output_name="output")
loop_layer = builder_top.add_loop("loop_layer")
loop_layer.loop.maxLoopIterations = 5
builder_body = NeuralNetworkBuilder(
input_features=None,
output_features=None,
spec=None,
nn_spec=loop_layer.loop.bodyNetwork,
)
builder_body.add_elementwise(
"add", input_names=["output"], output_name="x", mode="ADD", alpha=2
)
builder_body.add_copy("copy_2", input_name="x", output_name="output")
coremltools.models.utils.save_spec(
builder_top.spec, "/tmp/simple_loop_fixed_iterations.mlmodel"
)
mlmodel = MLModel(builder_top.spec)
# True branch case
input_dict = {"data": np.array([0], dtype="float")}
output_ref = {"output": np.array([10], dtype="float")}
self._test_model(mlmodel, input_dict, output_ref)
def test_simple_loop_fixed_iterations(self):
input_features = [('data', datatypes.Array(1))]
output_features = [('output', None)]
builder_top = NeuralNetworkBuilder(input_features,
output_features,
disable_rank5_shape_mapping=True)
builder_top.add_copy('copy_1', input_name='data', output_name='output')
loop_layer = builder_top.add_loop('loop_layer')
loop_layer.loop.maxLoopIterations = 5
builder_body = NeuralNetworkBuilder(
input_features=None,
output_features=None,
spec=None,
nn_spec=loop_layer.loop.bodyNetwork)
builder_body.add_elementwise('add',
input_names=['output'],
output_name='x',
mode='ADD',
alpha=2)
builder_body.add_copy('copy_2', input_name='x', output_name='output')
coremltools.models.utils.save_spec(
builder_top.spec, '/tmp/simple_loop_fixed_iterations.mlmodel')
mlmodel = MLModel(builder_top.spec)
# True branch case
input_dict = {'data': np.array([0], dtype='float')}
output_ref = {'output': np.array([10], dtype='float')}
self._test_model(mlmodel, input_dict, output_ref)
