def test_rebuild_submodel(model_2):
output_nodes = []
for output in model_2.outputs:
layer, node_index, tensor_index = output._keras_history
output_nodes.append(get_inbound_nodes(layer)[node_index])
surgeon = Surgeon(model_2)
outputs, _ = surgeon._rebuild_graph(model_2.inputs, output_nodes)
new_model = Model(model_2.inputs, outputs)
assert compare_models(model_2, new_model)
def test_rebuild_submodel(model_2):
output_nodes = [
model_2.output_layers[i].inbound_nodes[node_index]
for i, node_index in enumerate(model_2.output_layers_node_indices)
]
surgeon = Surgeon(model_2)
outputs, _ = surgeon._rebuild_graph(model_2.inputs, output_nodes)
new_model = Model(model_2.inputs, outputs)
assert compare_models(model_2, new_model)
