316 lines
11 KiB
Python
316 lines
11 KiB
Python
|
# Copyright 2020 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.
|
||
|
|
# ==============================================================================
|
||
|
|
"""Tests for saving/loading function for keras Model."""
|
||
|
|
|
||
|
|
import collections
|
||
|
|
|
||
|
|
import keras
|
||
|
|
|
||
|
|
# Declaring namedtuple()
|
||
|
|
ModelFn = collections.namedtuple(
|
||
|
|
"ModelFn", ["model", "input_shape", "target_shape"]
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def basic_sequential():
|
||
|
|
"""Basic sequential model."""
|
||
|
|
model = keras.Sequential(
|
||
|
|
[
|
||
|
|
keras.layers.Dense(3, activation="relu", input_shape=(3,)),
|
||
|
|
keras.layers.Dense(2, activation="softmax"),
|
||
|
|
]
|
||
|
|
)
|
||
|
|
return ModelFn(model, (None, 3), (None, 2))
|
||
|
|
|
||
|
|
|
||
|
|
def basic_sequential_deferred():
|
||
|
|
"""Sequential model with deferred input shape."""
|
||
|
|
model = keras.Sequential(
|
||
|
|
[
|
||
|
|
keras.layers.Dense(3, activation="relu"),
|
||
|
|
keras.layers.Dense(2, activation="softmax"),
|
||
|
|
]
|
||
|
|
)
|
||
|
|
return ModelFn(model, (None, 3), (None, 2))
|
||
|
|
|
||
|
|
|
||
|
|
def stacked_rnn():
|
||
|
|
"""Stacked RNN model."""
|
||
|
|
inputs = keras.Input((None, 3))
|
||
|
|
layer = keras.layers.RNN([keras.layers.LSTMCell(2) for _ in range(3)])
|
||
|
|
x = layer(inputs)
|
||
|
|
outputs = keras.layers.Dense(2)(x)
|
||
|
|
model = keras.Model(inputs, outputs)
|
||
|
|
return ModelFn(model, (None, 4, 3), (None, 2))
|
||
|
|
|
||
|
|
|
||
|
|
def lstm():
|
||
|
|
"""LSTM model."""
|
||
|
|
inputs = keras.Input((None, 3))
|
||
|
|
x = keras.layers.LSTM(4, return_sequences=True)(inputs)
|
||
|
|
x = keras.layers.LSTM(3, return_sequences=True)(x)
|
||
|
|
x = keras.layers.LSTM(2, return_sequences=False)(x)
|
||
|
|
outputs = keras.layers.Dense(2)(x)
|
||
|
|
model = keras.Model(inputs, outputs)
|
||
|
|
return ModelFn(model, (None, 4, 3), (None, 2))
|
||
|
|
|
||
|
|
|
||
|
|
def multi_input_multi_output():
|
||
|
|
"""Multi-input Multi-output model."""
|
||
|
|
body_input = keras.Input(shape=(None,), name="body")
|
||
|
|
tags_input = keras.Input(shape=(2,), name="tags")
|
||
|
|
|
||
|
|
x = keras.layers.Embedding(10, 4)(body_input)
|
||
|
|
body_features = keras.layers.LSTM(5)(x)
|
||
|
|
x = keras.layers.concatenate([body_features, tags_input])
|
||
|
|
|
||
|
|
pred_1 = keras.layers.Dense(2, activation="sigmoid", name="priority")(x)
|
||
|
|
pred_2 = keras.layers.Dense(3, activation="softmax", name="department")(x)
|
||
|
|
|
||
|
|
model = keras.Model(
|
||
|
|
inputs=[body_input, tags_input], outputs=[pred_1, pred_2]
|
||
|
|
)
|
||
|
|
return ModelFn(model, [(None, 1), (None, 2)], [(None, 2), (None, 3)])
|
||
|
|
|
||
|
|
|
||
|
|
def nested_sequential_in_functional():
|
||
|
|
"""A sequential model nested in a functional model."""
|
||
|
|
inner_model = keras.Sequential(
|
||
|
|
[
|
||
|
|
keras.layers.Dense(3, activation="relu", input_shape=(3,)),
|
||
|
|
keras.layers.Dense(2, activation="relu"),
|
||
|
|
]
|
||
|
|
)
|
||
|
|
|
||
|
|
inputs = keras.Input(shape=(3,))
|
||
|
|
x = inner_model(inputs)
|
||
|
|
outputs = keras.layers.Dense(2, activation="softmax")(x)
|
||
|
|
model = keras.Model(inputs, outputs)
|
||
|
|
return ModelFn(model, (None, 3), (None, 2))
|
||
|
|
|
||
|
|
|
||
|
|
def seq_to_seq():
|
||
|
|
"""Sequence to sequence model."""
|
||
|
|
num_encoder_tokens = 3
|
||
|
|
num_decoder_tokens = 3
|
||
|
|
latent_dim = 2
|
||
|
|
encoder_inputs = keras.Input(shape=(None, num_encoder_tokens))
|
||
|
|
encoder = keras.layers.LSTM(latent_dim, return_state=True)
|
||
|
|
_, state_h, state_c = encoder(encoder_inputs)
|
||
|
|
encoder_states = [state_h, state_c]
|
||
|
|
decoder_inputs = keras.Input(shape=(None, num_decoder_tokens))
|
||
|
|
decoder_lstm = keras.layers.LSTM(
|
||
|
|
latent_dim, return_sequences=True, return_state=True
|
||
|
|
)
|
||
|
|
decoder_outputs, _, _ = decoder_lstm(
|
||
|
|
decoder_inputs, initial_state=encoder_states
|
||
|
|
)
|
||
|
|
decoder_dense = keras.layers.Dense(num_decoder_tokens, activation="softmax")
|
||
|
|
decoder_outputs = decoder_dense(decoder_outputs)
|
||
|
|
model = keras.Model([encoder_inputs, decoder_inputs], decoder_outputs)
|
||
|
|
return ModelFn(
|
||
|
|
model,
|
||
|
|
[(None, 2, num_encoder_tokens), (None, 2, num_decoder_tokens)],
|
||
|
|
(None, 2, num_decoder_tokens),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def shared_layer_functional():
|
||
|
|
"""Shared layer in a functional model."""
|
||
|
|
main_input = keras.Input(shape=(10,), dtype="int32", name="main_input")
|
||
|
|
x = keras.layers.Embedding(output_dim=5, input_dim=4, input_length=10)(
|
||
|
|
main_input
|
||
|
|
)
|
||
|
|
lstm_out = keras.layers.LSTM(3)(x)
|
||
|
|
auxiliary_output = keras.layers.Dense(
|
||
|
|
1, activation="sigmoid", name="aux_output"
|
||
|
|
)(lstm_out)
|
||
|
|
auxiliary_input = keras.Input(shape=(5,), name="aux_input")
|
||
|
|
x = keras.layers.concatenate([lstm_out, auxiliary_input])
|
||
|
|
x = keras.layers.Dense(2, activation="relu")(x)
|
||
|
|
main_output = keras.layers.Dense(
|
||
|
|
1, activation="sigmoid", name="main_output"
|
||
|
|
)(x)
|
||
|
|
model = keras.Model(
|
||
|
|
inputs=[main_input, auxiliary_input],
|
||
|
|
outputs=[main_output, auxiliary_output],
|
||
|
|
)
|
||
|
|
return ModelFn(model, [(None, 10), (None, 5)], [(None, 1), (None, 1)])
|
||
|
|
|
||
|
|
|
||
|
|
def shared_sequential():
|
||
|
|
"""Shared sequential model in a functional model."""
|
||
|
|
inner_model = keras.Sequential(
|
||
|
|
[
|
||
|
|
keras.layers.Conv2D(2, 3, activation="relu"),
|
||
|
|
keras.layers.Conv2D(2, 3, activation="relu"),
|
||
|
|
]
|
||
|
|
)
|
||
|
|
inputs_1 = keras.Input((5, 5, 3))
|
||
|
|
inputs_2 = keras.Input((5, 5, 3))
|
||
|
|
x1 = inner_model(inputs_1)
|
||
|
|
x2 = inner_model(inputs_2)
|
||
|
|
x = keras.layers.concatenate([x1, x2])
|
||
|
|
outputs = keras.layers.GlobalAveragePooling2D()(x)
|
||
|
|
model = keras.Model([inputs_1, inputs_2], outputs)
|
||
|
|
return ModelFn(model, [(None, 5, 5, 3), (None, 5, 5, 3)], (None, 4))
|
||
|
|
|
||
|
|
|
||
|
|
class MySubclassModel(keras.Model):
|
||
|
|
"""A subclass model."""
|
||
|
|
|
||
|
|
def __init__(self, input_dim=3):
|
||
|
|
super().__init__(name="my_subclass_model")
|
||
|
|
self._config = {"input_dim": input_dim}
|
||
|
|
self.dense1 = keras.layers.Dense(8, activation="relu")
|
||
|
|
self.dense2 = keras.layers.Dense(2, activation="softmax")
|
||
|
|
self.bn = keras.layers.BatchNormalization()
|
||
|
|
self.dp = keras.layers.Dropout(0.5)
|
||
|
|
|
||
|
|
def call(self, inputs, **kwargs):
|
||
|
|
x = self.dense1(inputs)
|
||
|
|
x = self.dp(x)
|
||
|
|
x = self.bn(x)
|
||
|
|
return self.dense2(x)
|
||
|
|
|
||
|
|
def get_config(self):
|
||
|
|
return self._config
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def from_config(cls, config):
|
||
|
|
return cls(**config)
|
||
|
|
|
||
|
|
|
||
|
|
def nested_subclassed_model():
|
||
|
|
"""A subclass model nested in another subclass model."""
|
||
|
|
|
||
|
|
class NestedSubclassModel(keras.Model):
|
||
|
|
"""A nested subclass model."""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__()
|
||
|
|
self.dense1 = keras.layers.Dense(4, activation="relu")
|
||
|
|
self.dense2 = keras.layers.Dense(2, activation="relu")
|
||
|
|
self.bn = keras.layers.BatchNormalization()
|
||
|
|
self.inner_subclass_model = MySubclassModel()
|
||
|
|
|
||
|
|
def call(self, inputs):
|
||
|
|
x = self.dense1(inputs)
|
||
|
|
x = self.bn(x)
|
||
|
|
x = self.inner_subclass_model(x)
|
||
|
|
return self.dense2(x)
|
||
|
|
|
||
|
|
return ModelFn(NestedSubclassModel(), (None, 3), (None, 2))
|
||
|
|
|
||
|
|
|
||
|
|
def nested_subclassed_in_functional_model():
|
||
|
|
"""A subclass model nested in a functional model."""
|
||
|
|
inner_subclass_model = MySubclassModel()
|
||
|
|
inputs = keras.Input(shape=(3,))
|
||
|
|
x = inner_subclass_model(inputs)
|
||
|
|
x = keras.layers.BatchNormalization()(x)
|
||
|
|
outputs = keras.layers.Dense(2, activation="softmax")(x)
|
||
|
|
model = keras.Model(inputs, outputs)
|
||
|
|
return ModelFn(model, (None, 3), (None, 2))
|
||
|
|
|
||
|
|
|
||
|
|
def nested_functional_in_subclassed_model():
|
||
|
|
"""A functional model nested in a subclass model."""
|
||
|
|
|
||
|
|
def get_functional_model():
|
||
|
|
inputs = keras.Input(shape=(4,))
|
||
|
|
x = keras.layers.Dense(4, activation="relu")(inputs)
|
||
|
|
x = keras.layers.BatchNormalization()(x)
|
||
|
|
outputs = keras.layers.Dense(2)(x)
|
||
|
|
return keras.Model(inputs, outputs)
|
||
|
|
|
||
|
|
class NestedFunctionalInSubclassModel(keras.Model):
|
||
|
|
"""A functional nested in subclass model."""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__(name="nested_functional_in_subclassed_model")
|
||
|
|
self.dense1 = keras.layers.Dense(4, activation="relu")
|
||
|
|
self.dense2 = keras.layers.Dense(2, activation="relu")
|
||
|
|
self.inner_functional_model = get_functional_model()
|
||
|
|
|
||
|
|
def call(self, inputs):
|
||
|
|
x = self.dense1(inputs)
|
||
|
|
x = self.inner_functional_model(x)
|
||
|
|
return self.dense2(x)
|
||
|
|
|
||
|
|
return ModelFn(NestedFunctionalInSubclassModel(), (None, 3), (None, 2))
|
||
|
|
|
||
|
|
|
||
|
|
def shared_layer_subclassed_model():
|
||
|
|
"""Shared layer in a subclass model."""
|
||
|
|
|
||
|
|
class SharedLayerSubclassModel(keras.Model):
|
||
|
|
"""A subclass model with shared layers."""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__(name="shared_layer_subclass_model")
|
||
|
|
self.dense = keras.layers.Dense(3, activation="relu")
|
||
|
|
self.dp = keras.layers.Dropout(0.5)
|
||
|
|
self.bn = keras.layers.BatchNormalization()
|
||
|
|
|
||
|
|
def call(self, inputs):
|
||
|
|
x = self.dense(inputs)
|
||
|
|
x = self.dp(x)
|
||
|
|
x = self.bn(x)
|
||
|
|
return self.dense(x)
|
||
|
|
|
||
|
|
return ModelFn(SharedLayerSubclassModel(), (None, 3), (None, 3))
|
||
|
|
|
||
|
|
|
||
|
|
def functional_with_keyword_args():
|
||
|
|
"""A functional model with keyword args."""
|
||
|
|
inputs = keras.Input(shape=(3,))
|
||
|
|
x = keras.layers.Dense(4)(inputs)
|
||
|
|
x = keras.layers.BatchNormalization()(x)
|
||
|
|
outputs = keras.layers.Dense(2)(x)
|
||
|
|
|
||
|
|
model = keras.Model(inputs, outputs, name="m", trainable=False)
|
||
|
|
return ModelFn(model, (None, 3), (None, 2))
|
||
|
|
|
||
|
|
|
||
|
|
ALL_MODELS = [
|
||
|
|
("basic_sequential", basic_sequential),
|
||
|
|
("basic_sequential_deferred", basic_sequential_deferred),
|
||
|
|
("stacked_rnn", stacked_rnn),
|
||
|
|
("lstm", lstm),
|
||
|
|
("multi_input_multi_output", multi_input_multi_output),
|
||
|
|
("nested_sequential_in_functional", nested_sequential_in_functional),
|
||
|
|
("seq_to_seq", seq_to_seq),
|
||
|
|
("shared_layer_functional", shared_layer_functional),
|
||
|
|
("shared_sequential", shared_sequential),
|
||
|
|
("nested_subclassed_model", nested_subclassed_model),
|
||
|
|
(
|
||
|
|
"nested_subclassed_in_functional_model",
|
||
|
|
nested_subclassed_in_functional_model,
|
||
|
|
),
|
||
|
|
(
|
||
|
|
"nested_functional_in_subclassed_model",
|
||
|
|
nested_functional_in_subclassed_model,
|
||
|
|
),
|
||
|
|
("shared_layer_subclassed_model", shared_layer_subclassed_model),
|
||
|
|
("functional_with_keyword_args", functional_with_keyword_args),
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def get_models(exclude_models=None):
|
||
|
|
"""Get all models excluding the specified ones."""
|
||
|
|
models = [model for model in ALL_MODELS if model[0] not in exclude_models]
|
||
|
|
return models
|