You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
88 lines
3.3 KiB
88 lines
3.3 KiB
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
|
|
"""Utility functions for manipulating Keras models."""
|
|
|
|
import tensorflow as tf
|
|
|
|
|
|
def extract_submodel(model, inputs, outputs, name=None):
|
|
"""Extracts a section of a Keras model into a new model.
|
|
|
|
This method walks an existing model from the specified outputs back to the
|
|
specified inputs in order to construct a new model containing only a portion
|
|
of the old model, while sharing the layers and weights with the original
|
|
model.
|
|
|
|
WARNING: This method does not work for submodels containing layers that have
|
|
been used multiple times in the original model, or in other models beyond
|
|
the original model. (E.g. does not work for submodels that contain layers that
|
|
use shared weights). This also means that multiple overlapping submodels
|
|
cannot be extracted from the same model.
|
|
|
|
It also relies on recursion and will hit python's recursion limit for large
|
|
submodels.
|
|
|
|
Args:
|
|
model: The existing Keras model this method extracts a submodel from.
|
|
inputs: The layer inputs in the existing model that start the submodel
|
|
outputs: The layer outputs in the existing model that should be output by
|
|
the submodel
|
|
name: The name for the extracted model
|
|
|
|
Returns:
|
|
The extracted submodel specified by the given inputs and outputs
|
|
"""
|
|
output_to_layer = {}
|
|
output_to_layer_input = {}
|
|
for layer in model.layers:
|
|
layer_output = layer.output
|
|
layer_inputs = layer.input
|
|
output_to_layer[layer_output] = layer
|
|
output_to_layer_input[layer_output] = layer_inputs
|
|
|
|
model_inputs_dict = {}
|
|
memoized_results = {}
|
|
|
|
# Relies on recursion, very low limit in python
|
|
def _recurse_in_model(tensor):
|
|
"""Walk the existing model recursively to copy a submodel."""
|
|
if tensor in memoized_results:
|
|
return memoized_results[tensor]
|
|
if (tensor == inputs) or (isinstance(inputs, list) and tensor in inputs):
|
|
if tensor not in model_inputs_dict:
|
|
model_inputs_dict[tensor] = tf.keras.layers.Input(tensor=tensor)
|
|
out = model_inputs_dict[tensor]
|
|
else:
|
|
cur_inputs = output_to_layer_input[tensor]
|
|
cur_layer = output_to_layer[tensor]
|
|
if isinstance(cur_inputs, list):
|
|
out = cur_layer([_recurse_in_model(inp) for inp in cur_inputs])
|
|
else:
|
|
out = cur_layer(_recurse_in_model(cur_inputs))
|
|
memoized_results[tensor] = out
|
|
return out
|
|
|
|
if isinstance(outputs, list):
|
|
model_outputs = [_recurse_in_model(tensor) for tensor in outputs]
|
|
else:
|
|
model_outputs = _recurse_in_model(outputs)
|
|
|
|
if isinstance(inputs, list):
|
|
model_inputs = [model_inputs_dict[tensor] for tensor in inputs]
|
|
else:
|
|
model_inputs = model_inputs_dict[inputs]
|
|
|
|
return tf.keras.Model(inputs=model_inputs, outputs=model_outputs, name=name)
|