def _strip_clustering_wrapper(layer):
if isinstance(layer, cluster_wrapper.ClusterWeights):
if not hasattr(layer.layer, '_batch_input_shape') and\
hasattr(layer, '_batch_input_shape'):
layer.layer._batch_input_shape = layer._batch_input_shape
# We reset both arrays of weights, so that we can guarantee the correct
# order of newly created weights
layer.layer._trainable_weights = []
layer.layer._non_trainable_weights = []
for i in range(len(layer.restore)):
# This is why we used integers as keys
name, weight = layer.restore[i]
# In both cases we use k.batch_get_value since we need physical copies
# of the arrays to initialize a new tensor
if i in layer.gone_variables:
# If the variable was removed because it was clustered, we restore it
# by using updater we created earlier
new_weight_value = k.batch_get_value([weight()])[0]
else:
# If the value was not clustered(e.g. bias), we still store a valid
# reference to the tensor. We use this reference to get the value
new_weight_value = k.batch_get_value([weight])[0]
layer.layer.add_weight(
name=name,
shape=new_weight_value.shape,
initializer=initializers.Constant(new_weight_value),
trainable=True)
# When all weights are filled with the values, just return the underlying
# layer since it is now fully autonomous from its wrapper
return layer.layer
return layer
def F(inputs):
self.count += 1
R = fast_train_function(inputs)
if self.count % self.k == 0:
K.batch_get_value(slow_updates)
K.batch_get_value(copy_updates)
return R
def on_train_batch_end(self, batch, logs=None):
self.count += 1
if self.slow_weights is None:
with tf.control_dependencies(self.model.trainable_weights):
self.slow_weights = []
for fast_param in self.model.trainable_weights:
with ops.control_dependencies([fast_param]):
slow_param = tf.Variable(fast_param.initialized_value(),
dtype=fast_param.dtype,
trainable=False,
name=fast_param.name.split(":")[0])
self.slow_weights.append(slow_param)
K.track_variable(slow_param)
else:
if self.count % self.k == 0:
slow_ups, fast_ups = [], []
for fast, slow in zip(self.model.trainable_weights,
self.slow_weights):
slow_ups.append(K.update(slow, slow + self.alpha * (fast - slow)))
with tf.control_dependencies(slow_ups):
for fast, slow in zip(self.model.trainable_weights,
self.slow_weights):
fast_ups.append(K.update(fast, slow))
K.batch_get_value(slow_ups)
K.batch_get_value(fast_ups)
def _assert_weights_equal_value(self, annotated_weights, emulated_weights):
annotated_weight_values = K.batch_get_value(annotated_weights)
emulated_weight_values = K.batch_get_value(emulated_weights)
self.assertEqual(len(annotated_weight_values),
len(emulated_weight_values))
for aw, ew in zip(annotated_weight_values, emulated_weight_values):
self.assertAllClose(aw, ew)
def test_save_load_with_dense_features(self, tmpdir, api, loss, optimizer,
metrics):
if optimizer is None:
pytest.skip()
cols = [
feature_column_lib.numeric_column("a"),
feature_column_lib.indicator_column(
feature_column_lib.categorical_column_with_vocabulary_list(
"b", ["one", "two"])),
]
input_layers = {
"a": keras.layers.Input(shape=(1, ), name="a"),
"b": keras.layers.Input(shape=(1, ), name="b", dtype="string"),
}
fc_layer = dense_features.DenseFeatures(cols)(input_layers)
output = keras.layers.Dense(10)(fc_layer)
model = keras.models.Model(input_layers, output)
model.compile(
loss=loss,
optimizer=optimizer,
metrics=[metrics],
)
tiledb_uri = os.path.join(tmpdir, "model_array")
tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri,
model=model)
tiledb_model_obj.save(include_optimizer=True)
loaded_model = tiledb_model_obj.load(compile_model=True)
model_opt_weights = batch_get_value(getattr(model.optimizer,
"weights"))
loaded_opt_weights = batch_get_value(
getattr(loaded_model.optimizer, "weights"))
# Assert optimizer weights are equal
for weight_model, weight_loaded_model in zip(model_opt_weights,
loaded_opt_weights):
np.testing.assert_array_equal(weight_model, weight_loaded_model)
inputs_a = np.arange(10).reshape(10, 1)
inputs_b = np.arange(10).reshape(10, 1).astype("str")
# Assert model predictions are equal
np.testing.assert_array_equal(
loaded_model.predict({
"a": inputs_a,
"b": inputs_b
}),
model.predict({
"a": inputs_a,
"b": inputs_b
}),
)
def test_get_and_set_weights_does_not_add_ops(self):
table_handler = self.get_table_handler()
table_data = {b'a': 1, b'b': 2, b'c': 3}
table_handler.set_weights(
[list(table_data.keys()),
list(table_data.values())])
_ = backend.batch_get_value(table_handler.get_tensors())
backend.get_session().graph.finalize()
table_handler.set_weights(
[list(table_data.keys()),
list(table_data.values())])
_ = backend.batch_get_value(table_handler.get_tensors())
def on_epoch_end(self, batch, logs=None):
# At the end of every epoch, remask the weights. This ensures that when
# the model is saved after completion, the weights represent mask*weights.
layers = self.model.layers
weight_mask_ops = []
for layer in layers:
if isinstance(layer, pruning_wrapper.PruneLowMagnitude):
if tf.executing_eagerly():
layer.pruning_obj.weight_mask_op()
else:
weight_mask_ops.append(layer.pruning_obj.weight_mask_op())
K.batch_get_value(weight_mask_ops)
def get_weights(self):
"""Returns the current weights of the optimizer.
The weights of an optimizer are its state (ie, variables).
This function returns the weight values associated with this
optimizer as a list of Numpy arrays. The first value is always the
iterations count of the optimizer, followed by the optimizer's state
variables in the order they were created. The returned list can in turn
be used to load state into similarly parameterized optimizers.
For example, the RMSprop optimizer for this simple model returns a list of
three values-- the iteration count, followed by the root-mean-square value
of the kernel and bias of the single Dense layer:
>>> opt = tf.keras.optimizers.RMSprop()
>>> m = tf.keras.models.Sequential([tf.keras.layers.Dense(10)])
>>> m.compile(opt, loss='mse')
>>> data = np.arange(100).reshape(5, 20)
>>> labels = np.zeros(5)
>>> print('Training'); results = m.fit(data, labels)
Training ...
>>> len(opt.get_weights())
3
Returns:
Weights values as a list of numpy arrays.
"""
params = self.weights
return backend.batch_get_value(params)
def save_weights_to_hdf5_group(f, layers):
"""Saves the weights of a list of layers to a HDF5 group.
Arguments:
f: HDF5 group.
layers: List of layer instances.
"""
from tensorflow.python.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top
save_attributes_to_hdf5_group(
f, 'layer_names', [layer.name.encode('utf8') for layer in layers])
f.attrs['backend'] = K.backend().encode('utf8')
f.attrs['keras_version'] = str(keras_version).encode('utf8')
for layer in layers:
g = f.create_group(layer.name)
weights = _legacy_weights(layer)
weight_values = K.batch_get_value(weights)
weight_names = [w.name.encode('utf8') for w in weights]
save_attributes_to_hdf5_group(g, 'weight_names', weight_names)
for name, val in zip(weight_names, weight_values):
param_dset = g.create_dataset(name, val.shape, dtype=val.dtype)
if not val.shape:
# scalar
param_dset[()] = val
else:
param_dset[:] = val
def save_weights_to_hdf5_group(f, layers):
"""Saves the weights of a list of layers to a HDF5 group.
Arguments:
f: HDF5 group.
layers: List of layer instances.
"""
from tensorflow.python.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top
save_attributes_to_hdf5_group(
f, 'layer_names', [layer.name.encode('utf8') for layer in layers])
f.attrs['backend'] = K.backend().encode('utf8')
f.attrs['keras_version'] = str(keras_version).encode('utf8')
for layer in layers:
g = f.create_group(layer.name)
symbolic_weights = layer.weights
weight_values = K.batch_get_value(symbolic_weights)
weight_names = []
for i, (w, val) in enumerate(zip(symbolic_weights, weight_values)):
if hasattr(w, 'name') and w.name:
name = str(w.name)
else:
name = 'param_' + str(i)
weight_names.append(name.encode('utf8'))
save_attributes_to_hdf5_group(g, 'weight_names', weight_names)
for name, val in zip(weight_names, weight_values):
param_dset = g.create_dataset(name, val.shape, dtype=val.dtype)
if not val.shape:
# scalar
param_dset[()] = val
else:
param_dset[:] = val
def set_weights(self, weights):
"""Sets the weights of the optimizer, from Numpy arrays.
Should only be called after computing the gradients
(otherwise the optimizer has no weights).
Arguments:
weights: a list of Numpy arrays. The number
of arrays and their shape must match
number of the dimensions of the weights
of the optimizer (i.e. it should match the
output of `get_weights`).
Raises:
ValueError: in case of incompatible weight shapes.
"""
params = self.weights
if len(params) != len(weights):
raise ValueError(
'Length of the specified weight list (' + str(len(weights)) +
') does not match the number of weights '
'of the optimizer (' + str(len(params)) + ')')
weight_value_tuples = []
param_values = K.batch_get_value(params)
for pv, p, w in zip(param_values, params, weights):
if pv.shape != w.shape:
raise ValueError(
'Optimizer weight shape ' + str(pv.shape) + ' not compatible with '
'provided weight shape ' + str(w.shape))
weight_value_tuples.append((p, w))
K.batch_set_value(weight_value_tuples)
def on_epoch_end(self, batch, logs=None):
super(PruningSummaries, self).on_epoch_end(batch, logs)
pruning_logs = {}
params = []
layers = self.model.layers
for layer in layers:
if isinstance(layer, pruning_wrapper.PruneLowMagnitude):
for _, mask, threshold in layer.pruning_vars:
params.append(mask)
params.append(threshold)
params.append(self.model.optimizer.iterations)
values = K.batch_get_value(params)
iteration = values[-1]
del values[-1]
del params[-1]
param_value_pairs = list(zip(params, values))
for mask, mask_value in param_value_pairs[::2]:
pruning_logs.update(
{mask.name + '/sparsity': 1 - np.mean(mask_value)})
for threshold, threshold_value in param_value_pairs[1::2]:
pruning_logs.update(
{threshold.name + '/threshold': threshold_value})
self._log_metrics(pruning_logs, '', iteration)
def set_weights(self, weights):
"""Sets the weights of the optimizer, from Numpy arrays.
Should only be called after computing the gradients
(otherwise the optimizer has no weights).
Arguments:
weights: a list of Numpy arrays. The number of arrays and their shape
must match number of the dimensions of the weights of the optimizer
(i.e. it should match the output of `get_weights`).
Raises:
ValueError: in case of incompatible weight shapes.
"""
params = self.weights
if len(params) != len(weights):
raise ValueError('Length of the specified weight list (' +
str(len(weights)) +
') does not match the number of weights '
'of the optimizer (' + str(len(params)) + ')')
weight_value_tuples = []
param_values = K.batch_get_value(params)
for pv, p, w in zip(param_values, params, weights):
if pv.shape != w.shape:
raise ValueError('Optimizer weight shape ' + str(pv.shape) +
' not compatible with '
'provided weight shape ' + str(w.shape))
weight_value_tuples.append((p, w))
K.batch_set_value(weight_value_tuples)
def get_weights(self):
"""Returns the current value of the weights of the optimizer.
Returns:
A list of numpy arrays.
"""
return K.batch_get_value(self.weights)
def save_weights_to_hdf5_group(f, layers):
"""Saves the weights of a list of layers to a HDF5 group.
Arguments:
f: HDF5 group.
layers: List of layer instances.
"""
from tensorflow.python.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top
save_attributes_to_hdf5_group(
f, 'layer_names', [layer.name.encode('utf8') for layer in layers])
f.attrs['backend'] = K.backend().encode('utf8')
f.attrs['keras_version'] = str(keras_version).encode('utf8')
for layer in layers:
g = f.create_group(layer.name)
weights = _legacy_weights(layer)
weight_values = K.batch_get_value(weights)
weight_names = [w.name.encode('utf8') for w in weights]
save_attributes_to_hdf5_group(g, 'weight_names', weight_names)
for name, val in zip(weight_names, weight_values):
param_dset = g.create_dataset(name, val.shape, dtype=val.dtype)
if not val.shape:
# scalar
param_dset[()] = val
else:
param_dset[:] = val
def save_weights_to_hdf5_group(f, layers):
"""Saves the weights of a list of layers to a HDF5 group.
Arguments:
f: HDF5 group.
layers: List of layer instances.
"""
from tensorflow.python.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top
save_attributes_to_hdf5_group(
f, 'layer_names', [layer.name.encode('utf8') for layer in layers])
f.attrs['backend'] = K.backend().encode('utf8')
f.attrs['keras_version'] = str(keras_version).encode('utf8')
for layer in layers:
g = f.create_group(layer.name)
symbolic_weights = layer.weights
weight_values = K.batch_get_value(symbolic_weights)
weight_names = []
for i, (w, val) in enumerate(zip(symbolic_weights, weight_values)):
if hasattr(w, 'name') and w.name:
name = str(w.name)
else:
name = 'param_' + str(i)
weight_names.append(name.encode('utf8'))
save_attributes_to_hdf5_group(g, 'weight_names', weight_names)
for name, val in zip(weight_names, weight_values):
param_dset = g.create_dataset(name, val.shape, dtype=val.dtype)
if not val.shape:
# scalar
param_dset[()] = val
else:
param_dset[:] = val
def get_weights(self):
"""Returns the current value of the weights of the optimizer.
Returns:
A list of numpy arrays.
"""
return K.batch_get_value(self.weights)
def test_save_model_to_tiledb_array_weights(self, tmpdir, api, loss,
optimizer, metrics):
model = (api(num_hidden=1, num_classes=2, input_dim=3)
if api != ConfigSubclassModel else api(
hidden_units=[16, 16, 10]))
tiledb_uri = os.path.join(tmpdir, "model_array")
# Compiles the model if optimizer is present
if optimizer:
model.compile(loss=loss, optimizer=optimizer, metrics=[metrics])
input_shape = tuple(np.random.randint(20, size=2))
if not model.built:
model.build(input_shape)
tiledb_model_obj = TensorflowKerasTileDBModel(uri=tiledb_uri,
model=model)
tiledb_model_obj.save(include_optimizer=True if optimizer else False)
loaded_model = (
tiledb_model_obj.load(compile_model=True if optimizer else False)
if api != ConfigSubclassModel else tiledb_model_obj.load(
compile_model=True if optimizer else False,
custom_objects={"ConfigSubclassModel": ConfigSubclassModel},
input_shape=input_shape,
))
data = np.random.rand(
100, input_shape[-1] if api == ConfigSubclassModel else 3)
if optimizer:
model_opt_weights = batch_get_value(
getattr(model.optimizer, "weights"))
loaded_opt_weights = batch_get_value(
getattr(loaded_model.optimizer, "weights"))
# Assert optimizer weights are equal
for weight_model, weight_loaded_model in zip(
model_opt_weights, loaded_opt_weights):
np.testing.assert_array_equal(weight_model,
weight_loaded_model)
# Assert model predictions are equal
np.testing.assert_array_equal(loaded_model.predict(data),
model.predict(data))
def _write_array(
self,
include_optimizer: bool,
serialized_weights: bytes,
serialized_optimizer_weights: bytes,
meta: Optional[Meta],
) -> None:
"""Write Tensorflow model to a TileDB array."""
assert self.model
# TODO: Change timestamp when issue in core is resolved
with tiledb.open(self.uri,
"w",
timestamp=current_milli_time(),
ctx=self.ctx) as tf_model_tiledb:
if isinstance(self.model, (Functional, Sequential)):
tf_model_tiledb[:] = {
"model_weights": np.array([serialized_weights]),
"optimizer_weights":
np.array([serialized_optimizer_weights]),
}
else:
# Insert weights and optimizer
layer_names = []
weight_names = []
weight_values = []
for layer in sorted(self.model.layers, key=attrgetter("name")):
weights = layer.trainable_weights + layer.non_trainable_weights
weight_values.append(
pickle.dumps(backend.batch_get_value(weights)))
weight_names.append(
pickle.dumps([w.name.encode("utf8") for w in weights]))
layer_names.append(layer.name)
tf_model_tiledb[:] = {
"layer_name":
np.array(layer_names),
"weight_values":
np.array(weight_values),
"weight_names":
np.array(weight_names),
# TODO (TeamML) Fix array scheme to avoid optimizer_weight repetition. Nullable
"optimizer_weights":
np.array([
serialized_optimizer_weights
for _ in range(len(self.model.layers))
]),
}
# Insert all model metadata
model_metadata = saving_utils.model_metadata(
self.model, include_optimizer)
for key, value in model_metadata.items():
tf_model_tiledb.meta[key] = json.dumps(
value, default=json_utils.get_json_type).encode("utf8")
self.update_model_metadata(array=tf_model_tiledb, meta=meta)
def get_weights(self):
"""Retrieves the weights of the model.
Returns:
A flat list of Numpy arrays.
"""
weights = []
for cell in self.cells:
if isinstance(cell, Layer):
weights += cell.weights
return K.batch_get_value(weights)
def test_get_and_set_weights(self):
table_handler = self.get_table_handler()
table_data = {b'a': 1, b'b': 2, b'c': 3}
table_handler.set_weights(
[list(table_data.keys()),
list(table_data.values())])
weights = backend.batch_get_value(table_handler.get_tensors())
weight_data = {key: value for key, value in zip(weights[0], weights[1])}
self.assertDictEqual(table_data, weight_data)
def beam_search(predictions: List[np.ndarray],
padding_token_id: int,
start_sentence_token_id: int,
end_sentence_token_id: int,
beam_width: int = 5,
beam_top_paths: int = 5):
"""
predictions: output from a softmax layer, y true labels
# TODO if time permits implement own beam search, TF is too slow
"""
print("{}: In beam search".format(time.strftime(time.strftime("%Y-%m-%d-%H-%M%S"))))
start_time = time.time()
beam_search_predictions_list = []
beam_search_probs_list = []
for pred in predictions:
top_path_prediction_tensors, probs = K.ctc_decode(
np.expand_dims(pred, 0),
(pred.shape[0],),
greedy=False,
beam_width=beam_width,
top_paths=beam_top_paths
)
beam_search_predictions_list.append(top_path_prediction_tensors)
beam_search_probs_list.append(probs)
# evaluate tensorflow graph
print("{}: Evaluating beam search TF graph".format(time.strftime(time.strftime("%Y-%m-%d-%H-%M%S"))))
beam_search_predictions_evaluated: List[np.ndarray] = K.batch_get_value(beam_search_predictions_list)
print("{} Cleaning beamsearch results".format(time.strftime(time.strftime("%Y-%m-%d-%H-%M%S"))))
best_predictions = [list(trim_pred(pred, padding_token_id,
start_sentence_token_id,
end_sentence_token_id) for pred in beam_search_single_result)
for beam_search_single_result in beam_search_predictions_evaluated]
del beam_search_predictions_evaluated # freeup much needed memory
top_paths_predictions: np.ndarray = K.batch_get_value(beam_search_probs_list)
best_predictions_probs = list(map(lambda pred: np.exp(pred[0]), top_paths_predictions))
del top_paths_predictions # freeup much needed memory
print("beam search ended for one iteration in {}ms".format(time.time() - start_time))
return best_predictions, best_predictions_probs
def test_save_load_with_compile_functional(self):
functional_model = testing_utils.get_small_functional_mlp(
num_hidden=1, num_classes=2, input_dim=3)
functional_model = add_optimizer(functional_model)
tiledb_uri = os.path.join(self.get_temp_dir(), "model_array")
tiledb_model_obj = TensorflowTileDB(uri=tiledb_uri)
tiledb_model_obj.save(model=functional_model, include_optimizer=True)
loaded_model = tiledb_model_obj.load(compile_model=True)
data = np.random.rand(100, 3)
model_opt_weights = batch_get_value(
getattr(functional_model.optimizer, "weights"))
loaded_opt_weights = batch_get_value(
getattr(loaded_model.optimizer, "weights"))
# Assert optimizer weights are equal
for weight_model, weight_loaded_model in zip(model_opt_weights,
loaded_opt_weights):
np.testing.assert_array_equal(weight_model, weight_loaded_model)
# Assert model predictions are equal
np.testing.assert_array_equal(loaded_model.predict(data),
functional_model.predict(data))
def set_weights(self, weights):
"""Set the weights of the optimizer.
The weights of an optimizer are its state (ie, variables).
This function takes the weight values associated with this
optimizer as a list of Numpy arrays. The first value is always the
iterations count of the optimizer, followed by the optimizer's state
variables in the order they are created. The passed values are used to set
the new state of the optimizer.
For example, the RMSprop optimizer for this simple model takes a list of
three values-- the iteration count, followed by the root-mean-square value
of the kernel and bias of the single Dense layer:
>>> opt = tf.keras.optimizers.RMSprop()
>>> m = tf.keras.models.Sequential([tf.keras.layers.Dense(10)])
>>> m.compile(opt, loss='mse')
>>> data = np.arange(100).reshape(5, 20)
>>> labels = np.zeros(5)
>>> print('Training'); results = m.fit(data, labels)
Training ...
>>> new_weights = [np.array(10), np.ones([20, 10]), np.zeros([10])]
>>> opt.set_weights(new_weights)
>>> opt.iterations
<tf.Variable 'RMSprop/iter:0' shape=() dtype=int64, numpy=10>
Arguments:
weights: weight values as a list of numpy arrays.
"""
params = self.weights
if len(params) != len(weights):
raise ValueError(
"You called `set_weights(weights)` on optimizer " +
self._name + " with a weight list of length " +
str(len(weights)) + ", but the optimizer was expecting " +
str(len(params)) + " weights. Provided weights: " +
str(weights)[:50] + "...")
if not params:
return
weight_value_tuples = []
param_values = backend.batch_get_value(params)
for pv, p, w in zip(param_values, params, weights):
if pv.shape != w.shape:
raise ValueError("Optimizer weight shape " + str(pv.shape) +
" not compatible with "
"provided weight shape " + str(w.shape))
weight_value_tuples.append((p, w))
backend.batch_set_value(weight_value_tuples)
def set_weights(self, weights):
params = self.weights
if len(params) != len(weights):
raise ValueError(
"You called `set_weights(weights)` on optimizer " + self._name +
" with a weight list of length " + str(len(weights)) +
", but the optimizer was expecting " + str(len(params)) +
" weights. Provided weights: " + str(weights)[:50] + "...")
if not params:
return
weight_value_tuples = []
param_values = backend.batch_get_value(params)
for pv, p, w in zip(param_values, params, weights):
if pv.shape != w.shape:
raise ValueError("Optimizer weight shape " + str(pv.shape) +
" not compatible with "
"provided weight shape " + str(w.shape))
weight_value_tuples.append((p, w))
backend.batch_set_value(weight_value_tuples)
def save_weights_to_hdf5_group(f, layers):
"""Saves the weights of a list of layers to a HDF5 group.
Arguments:
f: HDF5 group.
layers: List of layer instances.
"""
from tensorflow.python.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top
save_attributes_to_hdf5_group(
f, 'layer_names', [layer.name.encode('utf8') for layer in layers])
f.attrs['backend'] = K.backend().encode('utf8')
f.attrs['keras_version'] = str(keras_version).encode('utf8')
# On TPUs, modifying the graph between session.runs() triggers some expensive
# recompilation overhead. To avoid this, we build up the full set of tensors
# to save before fetching weights, thus only modifying the graph once.
layer_weights_dict = {}
for layer in layers:
layer_weights_dict[layer.name] = [
ops.convert_to_tensor(w) for w in layer.weights
]
for layer in layers:
g = f.create_group(layer.name)
symbolic_weights = layer_weights_dict[layer.name]
weight_values = K.batch_get_value(symbolic_weights)
weight_names = []
for i, (w, val) in enumerate(zip(symbolic_weights, weight_values)):
if hasattr(w, 'name') and w.name:
name = str(w.name)
else:
name = 'param_' + str(i)
weight_names.append(name.encode('utf8'))
save_attributes_to_hdf5_group(g, 'weight_names', weight_names)
for name, val in zip(weight_names, weight_values):
param_dset = g.create_dataset(name, val.shape, dtype=val.dtype)
if not val.shape:
# scalar
param_dset[()] = val
else:
param_dset[:] = val
def train(self) -> Dict:
for data_set in self.data_sets:
data_set('train')
valid_exists = all(
data_set.valid_X is not None and data_set.valid_Y is not None
for data_set in self.data_sets)
self._model.fit(
**{
'x':
self.data_sets[0].data_X,
'y':
self.data_sets[0].data_Y,
'batch_size':
self.batch_size,
'validation_data': (
self.data_sets[0].valid_X,
self.data_sets[0].valid_Y) if valid_exists else None,
'epochs':
self.epochs,
'verbose':
0,
'callbacks': [
tf.keras.callbacks.ModelCheckpoint(save_weights_only=True,
save_best_only=True,
filepath=self.tmp_file,
verbose=0),
tf.keras.callbacks.
EarlyStopping(monitor='val_loss', patience=10)
] if valid_exists else None,
})
if valid_exists:
self._model.load_weights(self.tmp_file)
state = dict()
for f, (tfObj, vname_dict) in self._id2tfObj.items():
variables = tfObj.variables
state[f] = {
vname_dict[variable.name]: value
for variable, value in zip(variables,
K.batch_get_value(variables))
}
state[NeuralNetworkFrameworkInterface.arg_CMP] = self.cmp
return state
def save_tf_keras_optimizer(optimizer, h5py_file):
if isinstance(optimizer, optimizers.TFOptimizer):
logging.warning(
'TensorFlow optimizers do not '
'make it possible to access '
'optimizer attributes or optimizer state '
'after instantiation. '
'As a result, we cannot save the optimizer '
'as part of the model save file.'
'You will have to compile your model again after loading it. '
'Prefer using a Keras optimizer instead '
'(see keras.io/optimizers).')
else:
h5py_file.attrs['training_config'] = json.dumps(
{
'optimizer_config': {
'class_name': optimizer.__class__.__name__,
'config': optimizer.get_config()
}
},
default=serialization.get_json_type).encode('utf8')
# Save optimizer weights.
symbolic_weights = getattr(optimizer, 'weights')
if symbolic_weights:
optimizer_weights_group = h5py_file.create_group(
'optimizer_weights')
weight_values = K.batch_get_value(symbolic_weights)
weight_names = []
for w, val in zip(symbolic_weights, weight_values):
name = str(w.name)
weight_names.append(name.encode('utf8'))
optimizer_weights_group.attrs['weight_names'] = weight_names
for name, val in zip(weight_names, weight_values):
param_dset = optimizer_weights_group.create_dataset(
name, val.shape, dtype=val.dtype)
if not val.shape:
# scalar
param_dset[()] = val
else:
param_dset[:] = val
h5py_file.flush()
def save_weights_to_hdf5_group(f, layers):
"""Saves the weights of a list of layers to a HDF5 group.
Arguments:
f: HDF5 group.
layers: List of layer instances.
"""
from tensorflow.python.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top
save_attributes_to_hdf5_group(
f, 'layer_names', [layer.name.encode('utf8') for layer in layers])
f.attrs['backend'] = K.backend().encode('utf8')
f.attrs['keras_version'] = str(keras_version).encode('utf8')
# On TPUs, modifying the graph between session.runs() triggers some expensive
# recompilation overhead. To avoid this, we build up the full set of tensors
# to save before fetching weights, thus only modifying the graph once.
layer_weights_dict = {}
for layer in layers:
layer_weights_dict[layer.name] = [ops.convert_to_tensor(w)
for w in layer.weights]
for layer in layers:
g = f.create_group(layer.name)
symbolic_weights = layer_weights_dict[layer.name]
weight_values = K.batch_get_value(symbolic_weights)
weight_names = []
for i, (w, val) in enumerate(zip(symbolic_weights, weight_values)):
if hasattr(w, 'name') and w.name:
name = str(w.name)
else:
name = 'param_' + str(i)
weight_names.append(name.encode('utf8'))
save_attributes_to_hdf5_group(g, 'weight_names', weight_names)
for name, val in zip(weight_names, weight_values):
param_dset = g.create_dataset(name, val.shape, dtype=val.dtype)
if not val.shape:
# scalar
param_dset[()] = val
else:
param_dset[:] = val
def save_optimizer_weights_to_hdf5_group(hdf5_group, optimizer):
"""Saves optimizer weights of a optimizer to a HDF5 group.
Arguments:
hdf5_group: HDF5 group.
optimizer: optimizer instance.
"""
symbolic_weights = getattr(optimizer, 'weights')
if symbolic_weights:
weights_group = hdf5_group.create_group('optimizer_weights')
weight_names = [str(w.name).encode('utf8') for w in symbolic_weights]
save_attributes_to_hdf5_group(weights_group, 'weight_names', weight_names)
weight_values = K.batch_get_value(symbolic_weights)
for name, val in zip(weight_names, weight_values):
param_dset = weights_group.create_dataset(
name, val.shape, dtype=val.dtype)
if not val.shape:
# scalar
param_dset[()] = val
else:
param_dset[:] = val
def save_weights_to_hdf5_group(f, fh_dict, layers, compress=False):
"""Saves the weights of a list of layers to a HDF5 group.
This is revised version. We split the attributes of HDF5 group into another
JSON file to avoid the heading memory excessing problem. Compared to original
Keras API, we need to load an extra file IO handle, fh_dict.
Arguments:
f: HDF5 group.
fh_dict: JSON config dictionary.
layers: a list of layer instances.
compress: whether to compress the weights.
"""
compression = 'gzip' if compress else None
save_attributes_to_hdf5_group(fh_dict, f.name, 'layer_names',
[layer.name for layer in layers])
for layer in layers:
g = f.create_group(layer.name)
symbolic_weights = layer.weights
weight_values = K.batch_get_value(symbolic_weights)
weight_names = []
for i, (w, val) in enumerate(zip(symbolic_weights, weight_values)):
if hasattr(w, 'name') and w.name:
name = str(w.name)
else:
name = 'param_' + str(i)
weight_names.append(name)
save_attributes_to_hdf5_group(fh_dict, g.name, 'weight_names',
weight_names)
for name, val in zip(weight_names, weight_values):
param_dset = g.create_dataset(
name,
val.shape,
dtype=val.dtype,
compression=(compression if val.shape else None))
if not val.shape:
# scalar
param_dset[()] = val
else:
param_dset[:] = val
def save_optimizer_weights_to_hdf5_group(hdf5_group, optimizer):
"""Saves optimizer weights of a optimizer to a HDF5 group.
Arguments:
hdf5_group: HDF5 group.
optimizer: optimizer instance.
"""
symbolic_weights = getattr(optimizer, 'weights')
if symbolic_weights:
weights_group = hdf5_group.create_group('optimizer_weights')
weight_names = [str(w.name).encode('utf8') for w in symbolic_weights]
save_attributes_to_hdf5_group(weights_group, 'weight_names', weight_names)
weight_values = K.batch_get_value(symbolic_weights)
for name, val in zip(weight_names, weight_values):
param_dset = weights_group.create_dataset(
name, val.shape, dtype=val.dtype)
if not val.shape:
# scalar
param_dset[()] = val
else:
param_dset[:] = val
def get_weights(self): params = self.weights return backend.batch_get_value(params)
def experimental_tpu_predict_loop(model,
dataset,
verbose=0,
steps=None,
callbacks=None):
"""Predict loop for predicting with TPU DistributionStrategy.
Arguments:
model: Keras Model instance.
dataset: Dataset for input data.
verbose: Integer, Verbosity mode 0 or 1.
steps: Total number of steps (batches of samples)
before declaring `_predict_loop` finished.
Ignored with the default value of `None`.
callbacks: List of callbacks to be called during training
Returns:
Array of predictions (if the model has a single output)
or list of arrays of predictions
(if the model has multiple outputs).
"""
mode = ModeKeys.PREDICT
steps = training_utils.infer_steps_for_dataset(dataset, steps,
steps_name='steps')
dataset_fully_shaped = (distributed_training_utils.
is_dataset_shape_fully_defined(dataset))
padding_handler = None
if not dataset_fully_shaped:
# TODO(hongjunchoi): Investigate whether operations from
# PartialBatchPaddingHandler are unnecessarily pruned out
# during graph optimization.
padding_handler = padding_util.PartialBatchPaddingHandler(
model._feed_output_shapes)
batch_size, _, prefetch_buffer = input_lib._get_dataset_attributes(dataset)
padding_handler.padded_batch_size = batch_size
padding_handler.padding_mask = dataset.reduce(padding_handler.padding_mask,
padding_handler.update_mask)
dataset = dataset.map(padding_handler.pad_batch)
dataset = dataset.apply(batching.unbatch())
# Upon this point, it is guaranteed that the dataset does not
# have partial batches. Thus, we set `drop_remainder=True` to
# get static shape information about the elements in the dataset.
dataset = dataset.batch(batch_size, drop_remainder=True)
if prefetch_buffer is not None:
dataset = dataset.prefetch(prefetch_buffer)
current_strategy = model._distribution_strategy
iterator = distributed_training_utils.get_iterator(dataset, current_strategy)
scope = distributed_training_utils.distributed_scope(
strategy=current_strategy, learning_phase=0)
scope.__enter__()
out_labels = model.output_names
step_fn = _make_step_fn(model, ModeKeys.PREDICT, current_strategy, out_labels)
# Add initial dummy values for outputs.
initial_loop_values = {}
batch_dimension = distributed_training_utils.get_batch_dimension(iterator)
for name, tensor in zip(model.output_names, model.outputs):
# TODO(priyag): This is a workaround as we do not know the batch dimension
# of the model's output at this point.
shape = tensor_shape.TensorShape(tensor.shape.dims)
shape.dims = [batch_dimension] + shape.dims[1:]
initial_loop_values[name] = array_ops.zeros(shape, tensor.dtype)
# TODO(priyag, sourabhbajaj): Support steps_per_run if/when we add outfeed.
ctx = current_strategy.extended.experimental_run_steps_on_iterator(
step_fn, iterator, iterations=1,
initial_loop_values=initial_loop_values)
predict_op = ctx.run_op
output_tensors = ctx.last_step_outputs
if verbose == 1:
progbar = Progbar(target=steps)
if model._compile_distribution:
distributed_training_utils._copy_weights_to_distributed_model(model, mode)
distributed_training_utils._reset_metrics(model)
callbacks = cbks.configure_callbacks(
callbacks,
model,
do_validation=False,
epochs=1,
steps_per_epoch=steps,
verbose=verbose,
count_mode='steps',
mode=mode)
callbacks._call_begin_hook(mode)
# Since we do not know how many samples we will see, we cannot pre-allocate
# the returned Numpy arrays. Instead, we store one array per batch seen
# and concatenate them upon returning.
unconcatenated_outs = [[] for _ in model.outputs]
if steps is not None:
target_steps = steps
else:
target_steps = np.inf
current_step = 0
while current_step < target_steps:
batch_logs = {'batch': current_step, 'size': 1}
callbacks._call_batch_hook(mode, 'begin', current_step, batch_logs)
try:
_, batch_outs = K.batch_get_value([predict_op, output_tensors])
except errors.OutOfRangeError:
if steps is not None:
warning_msg = 'Make sure that your dataset can generate at least '
'`steps` batches (in this case, {} batches).'.format(steps)
else:
warning_msg = 'Number of steps ran: {} steps'.format(current_step)
logging.warning('Your dataset iterator ran out of data; '
'interrupting evaluation. ' + warning_msg)
break
# TODO(priyag): maybe need to unwrap the outputs first for MirroredStrategy.
for i, label in enumerate(model.output_names):
unconcatenated_outs[i].extend(batch_outs[label])
batch_logs = cbks.make_logs(model, batch_logs, batch_outs, mode)
callbacks._call_batch_hook(mode, 'end', current_step, batch_logs)
if verbose >= 1:
progbar.update(current_step + 1)
current_step += 1
callbacks._call_end_hook(mode)
scope.__exit__(None, None, None)
if len(unconcatenated_outs) == 1:
prediction_result = np.concatenate(unconcatenated_outs[0], axis=0)
else:
prediction_result = [
np.concatenate(unconcatenated_outs[i], axis=0)
for i in range(len(unconcatenated_outs))
]
if padding_handler:
prediction_result = padding_handler.apply_mask(prediction_result)
return prediction_result
def experimental_tpu_test_loop(model,
dataset,
verbose=0,
steps=None,
callbacks=None):
"""Test loop for evaluating with TPU DistributionStrategy.
Arguments:
model: Keras Model instance.
dataset: Dataset for input data.
verbose: Integer, Verbosity mode 0 or 1.
steps: Total number of steps (batches of samples)
before declaring predictions finished.
Ignored with the default value of `None`.
callbacks: List of callbacks to be called during training
Returns:
Scalar loss (if the model has a single output and no metrics)
or list of scalars (if the model has multiple outputs
and/or metrics). The attribute `model.metrics_names` will give you
the display labels for the outputs.
"""
mode = ModeKeys.TEST
current_strategy = model._distribution_strategy
iterator = distributed_training_utils.get_iterator(dataset,
current_strategy)
steps = training_utils.infer_steps_for_dataset(dataset, steps,
steps_name='steps')
scope = distributed_training_utils.distributed_scope(
strategy=current_strategy, learning_phase=0)
scope.__enter__()
out_labels = model.metrics_names
def _test_step_fn(inputs):
"""A fn that returns output of single test step."""
inputs, targets = inputs
(distribution_strategy_context.get_replica_context().merge_call(
_build_model, args=(model, mode, inputs, targets)))
(_, outputs, updates, _) = (
_per_device_execution_function(
distributed_training_utils.get_distributed_model(model, mode),
mode))
with ops.control_dependencies([updates]):
return outputs
test_input_data = iterator.get_next()
per_replica_outputs = current_strategy.experimental_run_v2(
_test_step_fn, args=(test_input_data,))
output_tensors = {}
for label, output in zip(out_labels, per_replica_outputs):
if label == 'loss':
reduce_op = ds_reduce_util.ReduceOp.SUM
else:
# We reduce all other metrics using mean for now. This is temporary
# workaround until new metrics are in place.
reduce_op = ds_reduce_util.ReduceOp.MEAN
output_tensors[label] = current_strategy.reduce(reduce_op, output)
test_op = control_flow_ops.group(list(output_tensors.values()))
if verbose >= 1:
progbar = Progbar(target=steps)
if model._compile_distribution:
distributed_training_utils._copy_weights_to_distributed_model(model, mode)
distributed_training_utils._reset_metrics(model)
callbacks = cbks.configure_callbacks(
callbacks,
model,
do_validation=False,
epochs=1,
steps_per_epoch=steps,
verbose=verbose,
count_mode='steps',
mode=ModeKeys.TEST)
callbacks._call_begin_hook(mode)
outs = [0.] * len(model.metrics_names)
if steps is not None:
target_steps = steps
else:
raise ValueError('Number of steps could not be infered from the data, '
'please pass the steps argument.')
current_step = 0
while current_step < target_steps:
batch_logs = {'batch': current_step, 'size': 1}
callbacks._call_batch_hook(mode, 'begin', current_step, batch_logs)
try:
_, batch_outs = K.batch_get_value([test_op, output_tensors])
except errors.OutOfRangeError:
warning_msg = 'Make sure that your dataset can generate at least '
'`steps` batches (in this case, {} batches).'.format(steps)
logging.warning('Your dataset iterator ran out of data; '
'interrupting evaluation. ' + warning_msg)
target_steps = current_step
break
for i, label in enumerate(model.metrics_names):
if i == 0:
# Loss is stateless metrics.
outs[i] += batch_outs[label]
else:
# For all stateful metrics, the aggregation is handled by mirrored vars.
outs[i] = batch_outs[label]
batch_logs = cbks.make_logs(model, batch_logs, outs, mode)
callbacks._call_batch_hook(mode, 'end', current_step, batch_logs)
if verbose == 1:
progbar.update(current_step + 1)
current_step += 1
if verbose >= 1:
# Progress bar finishes at the end.
progbar.update(target_steps)
callbacks._call_end_hook(mode)
scope.__exit__(None, None, None)
if len(outs) >= 0:
outs[0] /= (target_steps)
if len(outs) == 1:
return outs[0]
return outs
def experimental_tpu_predict_loop(model,
dataset,
verbose=0,
steps=None,
callbacks=None):
"""Predict loop for predicting with TPU DistributionStrategy.
Arguments:
model: Keras Model instance.
dataset: Dataset for input data.
verbose: Integer, Verbosity mode 0 or 1.
steps: Total number of steps (batches of samples)
before declaring `_predict_loop` finished.
Ignored with the default value of `None`.
callbacks: List of callbacks to be called during training
Returns:
Array of predictions (if the model has a single output)
or list of arrays of predictions
(if the model has multiple outputs).
"""
mode = ModeKeys.PREDICT
steps = training_utils.infer_steps_for_dataset(dataset, steps,
steps_name='steps')
dataset_fully_shaped = (distributed_training_utils.
is_dataset_shape_fully_defined(dataset))
padding_handler = None
if not dataset_fully_shaped:
# TODO(hongjunchoi): Investigate whether operations from
# PartialBatchPaddingHandler are unnecessarily pruned out
# during graph optimization.
padding_handler = padding_util.PartialBatchPaddingHandler(
model._feed_output_shapes)
batch_size, _, prefetch_buffer = input_lib._get_dataset_attributes(dataset)
padding_handler.padded_batch_size = batch_size
padding_handler.padding_mask = dataset.reduce(padding_handler.padding_mask,
padding_handler.update_mask)
dataset = dataset.map(padding_handler.pad_batch)
dataset = dataset.apply(batching.unbatch())
# Upon this point, it is guaranteed that the dataset does not
# have partial batches. Thus, we set `drop_remainder=True` to
# get static shape information about the elements in the dataset.
dataset = dataset.batch(batch_size, drop_remainder=True)
if prefetch_buffer is not None:
dataset = dataset.prefetch(prefetch_buffer)
current_strategy = model._distribution_strategy
iterator = distributed_training_utils.get_iterator(dataset, current_strategy)
scope = distributed_training_utils.distributed_scope(
strategy=current_strategy, learning_phase=0)
scope.__enter__()
def _predict_step_fn(inputs):
"""A fn that returns output of single prediction step."""
(distribution_strategy_context.get_replica_context().merge_call(
_build_model, args=(model, mode, inputs)))
(_, outputs, updates, _) = (
_per_device_execution_function(
distributed_training_utils.get_distributed_model(model, mode),
mode))
with ops.control_dependencies([updates]):
return outputs
# TODO(hongjunchoi): When numpy array is passed as an input to `predict()`
# use numpy arrays directly to avoid cumulating unnecessary input pipeline
# ops.
predict_input_data = iterator.get_next()
per_replica_outputs = current_strategy.experimental_run_v2(
_predict_step_fn, args=(predict_input_data,))
output_tensors = distributed_training_utils.flatten_perdevice_values(
current_strategy, per_replica_outputs)
if verbose >= 1:
progbar = Progbar(target=steps)
if model._compile_distribution:
distributed_training_utils._copy_weights_to_distributed_model(model, mode)
distributed_training_utils._reset_metrics(model)
callbacks = cbks.configure_callbacks(
callbacks,
model,
do_validation=False,
epochs=1,
steps_per_epoch=steps,
verbose=verbose,
count_mode='steps',
mode=mode)
callbacks._call_begin_hook(mode)
# Since we do not know how many samples we will see, we cannot pre-allocate
# the returned Numpy arrays. Instead, we store one array per batch seen
# and concatenate them upon returning.
num_model_outputs = len(model.output_names)
unconcatenated_outs = [[] for _ in range(num_model_outputs)]
if steps is not None:
target_steps = steps
else:
raise ValueError('Number of steps could not be infered from the data, '
'please pass the steps argument.')
current_step = 0
while current_step < target_steps:
batch_logs = {'batch': current_step, 'size': 1}
callbacks._call_batch_hook(mode, 'begin', current_step, batch_logs)
try:
predict_ops = control_flow_ops.group(output_tensors)
_, batch_outs = K.batch_get_value([predict_ops, output_tensors])
except errors.OutOfRangeError:
warning_msg = 'Make sure that your dataset can generate at least '
'`steps` batches (in this case, {} batches).'.format(steps)
logging.warning('Your dataset iterator ran out of data; '
'interrupting evaluation. ' + warning_msg)
break
# TODO(priyag): maybe need to unwrap the outputs first for MirroredStrategy.
for i in range(num_model_outputs):
output_start_index = i * current_strategy.num_replicas_in_sync
output_end_index = (
output_start_index + current_strategy.num_replicas_in_sync)
single_model_output = batch_outs[output_start_index:output_end_index]
unconcatenated_outs[i].extend(single_model_output)
batch_logs = cbks.make_logs(model, batch_logs, batch_outs, mode)
callbacks._call_batch_hook(mode, 'end', current_step, batch_logs)
if verbose == 1:
progbar.update(current_step + 1)
current_step += 1
if verbose >= 1:
# Progress bar finishes at the end.
progbar.update(current_step)
callbacks._call_end_hook(mode)
scope.__exit__(None, None, None)
if len(unconcatenated_outs) == 1:
prediction_result = np.concatenate(unconcatenated_outs[0], axis=0)
else:
prediction_result = [
np.concatenate(unconcatenated_outs[i], axis=0)
for i in range(len(unconcatenated_outs))
]
if padding_handler:
prediction_result = padding_handler.apply_mask(prediction_result)
return prediction_result
def experimental_tpu_test_loop(model,
dataset,
verbose=0,
steps=None,
callbacks=None):
"""Test loop for evaluating with TPU DistributionStrategy.
Arguments:
model: Keras Model instance.
dataset: Dataset for input data.
verbose: Integer, Verbosity mode 0 or 1.
steps: Total number of steps (batches of samples)
before declaring predictions finished.
Ignored with the default value of `None`.
callbacks: List of callbacks to be called during training
Returns:
Scalar loss (if the model has a single output and no metrics)
or list of scalars (if the model has multiple outputs
and/or metrics). The attribute `model.metrics_names` will give you
the display labels for the outputs.
"""
mode = ModeKeys.TEST
current_strategy = model._distribution_strategy
iterator = distributed_training_utils.get_iterator(dataset,
current_strategy)
steps = training_utils.infer_steps_for_dataset(dataset, steps,
steps_name='steps')
scope = distributed_training_utils.distributed_scope(
strategy=current_strategy, learning_phase=0)
scope.__enter__()
out_labels = model.metrics_names
step_fn = _make_step_fn(model, ModeKeys.TEST, current_strategy, out_labels)
# Add initial dummy values for loss and other metric tensors.
initial_loop_values = {}
initial_loop_values['loss'] = constant_op.constant(1e7)
for name in model.metrics_names[1:]:
tensor = model._all_stateful_metrics_tensors[name]
initial_loop_values[name] = array_ops.zeros(tensor.shape, tensor.dtype)
# TODO(priyag): Use steps_per_run when we use new metrics as they will
# allow handling metric computation at each step using variables.
ctx = current_strategy.extended.experimental_run_steps_on_iterator(
step_fn, iterator, iterations=1,
initial_loop_values=initial_loop_values)
test_op = ctx.run_op
output_tensors = ctx.last_step_outputs
if verbose == 1:
progbar = Progbar(target=steps)
if model._compile_distribution:
distributed_training_utils._copy_weights_to_distributed_model(model, mode)
distributed_training_utils._reset_metrics(model)
callbacks = cbks.configure_callbacks(
callbacks,
model,
do_validation=False,
epochs=1,
steps_per_epoch=steps,
verbose=verbose,
count_mode='steps',
mode=ModeKeys.TEST)
callbacks._call_begin_hook(mode)
outs = [0.] * len(model.metrics_names)
if steps is not None:
target_steps = steps
else:
target_steps = np.inf
current_step = 0
while current_step < target_steps:
batch_logs = {'batch': current_step, 'size': 1}
callbacks._call_batch_hook(mode, 'begin', current_step, batch_logs)
try:
_, batch_outs = K.batch_get_value([test_op, output_tensors])
except errors.OutOfRangeError:
if steps is not None:
warning_msg = 'Make sure that your dataset can generate at least '
'`steps` batches (in this case, {} batches).'.format(steps)
else:
warning_msg = 'Number of steps ran: {} steps'.format(current_step)
logging.warning('Your dataset iterator ran out of data; '
'interrupting evaluation. ' + warning_msg)
target_steps = current_step
break
for i, label in enumerate(model.metrics_names):
if i == 0:
# Loss is stateless metrics.
outs[i] += batch_outs[label]
else:
# For all stateful metrics, the aggregation is handled by mirrored vars.
outs[i] = batch_outs[label]
batch_logs = cbks.make_logs(model, batch_logs, outs, mode)
callbacks._call_batch_hook(mode, 'end', current_step, batch_logs)
if verbose >= 1:
progbar.update(current_step + 1)
current_step += 1
callbacks._call_end_hook(mode)
scope.__exit__(None, None, None)
if len(outs) >= 0:
outs[0] /= (target_steps)
if len(outs) == 1:
return outs[0]
return outs
def save_model(model, filepath, overwrite=True, include_optimizer=True):
"""Saves a model to a HDF5 file.
The saved model contains:
- the model's configuration (topology)
- the model's weights
- the model's optimizer's state (if any)
Thus the saved model can be reinstantiated in
the exact same state, without any of the code
used for model definition or training.
Arguments:
model: Keras model instance to be saved.
filepath: One of the following:
- String, path where to save the model
- `h5py.File` object where to save the model
overwrite: Whether we should overwrite any existing
model at the target location, or instead
ask the user with a manual prompt.
include_optimizer: If True, save optimizer's state together.
Raises:
ImportError: if h5py is not available.
"""
if h5py is None:
raise ImportError('`save_model` requires h5py.')
from tensorflow.python.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top
if not isinstance(filepath, h5py.File):
# If file exists and should not be overwritten.
if not overwrite and os.path.isfile(filepath):
proceed = ask_to_proceed_with_overwrite(filepath)
if not proceed:
return
f = h5py.File(filepath, mode='w')
opened_new_file = True
else:
f = filepath
opened_new_file = False
try:
f.attrs['keras_version'] = str(keras_version).encode('utf8')
f.attrs['backend'] = K.backend().encode('utf8')
f.attrs['model_config'] = json.dumps(
{
'class_name': model.__class__.__name__,
'config': model.get_config()
},
default=serialization.get_json_type).encode('utf8')
model_weights_group = f.create_group('model_weights')
model_layers = model.layers
save_weights_to_hdf5_group(model_weights_group, model_layers)
if include_optimizer and model.optimizer:
if isinstance(model.optimizer, optimizers.TFOptimizer):
logging.warning(
'TensorFlow optimizers do not '
'make it possible to access '
'optimizer attributes or optimizer state '
'after instantiation. '
'As a result, we cannot save the optimizer '
'as part of the model save file.'
'You will have to compile your model again after loading it. '
'Prefer using a Keras optimizer instead '
'(see keras.io/optimizers).')
else:
f.attrs['training_config'] = json.dumps(
{
'optimizer_config': {
'class_name': model.optimizer.__class__.__name__,
'config': model.optimizer.get_config()
},
'loss': model.loss,
'metrics': model.metrics,
'sample_weight_mode': model.sample_weight_mode,
'loss_weights': model.loss_weights,
},
default=serialization.get_json_type).encode('utf8')
# Save optimizer weights.
symbolic_weights = getattr(model.optimizer, 'weights')
if symbolic_weights:
optimizer_weights_group = f.create_group('optimizer_weights')
weight_values = K.batch_get_value(symbolic_weights)
weight_names = []
for w, val in zip(symbolic_weights, weight_values):
name = str(w.name)
weight_names.append(name.encode('utf8'))
optimizer_weights_group.attrs['weight_names'] = weight_names
for name, val in zip(weight_names, weight_values):
param_dset = optimizer_weights_group.create_dataset(
name, val.shape, dtype=val.dtype)
if not val.shape:
# scalar
param_dset[()] = val
else:
param_dset[:] = val
f.flush()
finally:
if opened_new_file:
f.close()
def experimental_tpu_fit_loop(model,
dataset,
epochs=100,
verbose=1,
callbacks=None,
initial_epoch=0,
steps_per_epoch=None,
val_dataset=None,
validation_steps=None,
validation_freq=1):
"""Fit loop for training with TPU DistributionStrategy.
Arguments:
model: Keras Model instance.
dataset: Dataset that returns inputs and targets
epochs: Number of times to iterate over the data
verbose: Integer, Verbosity mode, 0, 1 or 2
callbacks: List of callbacks to be called during training
initial_epoch: Epoch at which to start training
(useful for resuming a previous training run)
steps_per_epoch: Total number of steps (batches of samples)
before declaring one epoch finished and starting the
next epoch. Ignored with the default value of `None`.
val_dataset: Dataset for validation data.
validation_steps: Number of steps to run validation for
(only if doing validation from data tensors).
Ignored with the default value of `None`.
validation_freq: Only relevant if validation data is provided. Integer or
`collections.Container` instance (e.g. list, tuple, etc.). If an
integer, specifies how many training epochs to run before a new
validation run is performed, e.g. `validation_freq=2` runs
validation every 2 epochs. If a Container, specifies the epochs on
which to run validation, e.g. `validation_freq=[1, 2, 10]` runs
validation at the end of the 1st, 2nd, and 10th epochs.
Returns:
Returns `None`.
Raises:
ValueError: in case of invalid arguments.
"""
mode = ModeKeys.TRAIN
# TODO(fchollet): add support for `steps_per_epoch=None` in TPU loops.
current_strategy = model._distribution_strategy
iterator = distributed_training_utils.get_iterator(dataset, current_strategy)
steps_per_epoch = training_utils.infer_steps_for_dataset(
dataset, steps_per_epoch, epochs, steps_name='steps_per_epoch')
if (current_strategy.extended.steps_per_run != 1 and
steps_per_epoch is None):
raise ValueError('`steps_per_epoch` should be specified when calling '
'`fit` on the model with TPUStrategy when '
'`steps_per_run` != 1 .')
scope = distributed_training_utils.distributed_scope(
strategy=current_strategy, learning_phase=1)
scope.__enter__()
out_labels = model.metrics_names or []
step_fn = _make_step_fn(model, ModeKeys.TRAIN, current_strategy, out_labels)
# Add initial dummy values for loss and other metric tensors.
initial_loop_values = {}
initial_loop_values['loss'] = constant_op.constant(1e7)
for name in model.metrics_names[1:]:
tensor = model._all_stateful_metrics_tensors[name]
initial_loop_values[name] = array_ops.zeros(tensor.shape, tensor.dtype)
use_steps = steps_per_epoch is not None
if use_steps:
iteration_value = min(steps_per_epoch,
current_strategy.extended.steps_per_run)
else:
iteration_value = current_strategy.extended.steps_per_run
steps_per_run = K.variable(
value=iteration_value,
dtype='int32',
name='steps_per_run')
ctx = current_strategy.extended.experimental_run_steps_on_iterator(
step_fn, iterator, iterations=steps_per_run,
initial_loop_values=initial_loop_values)
train_op = ctx.run_op
output_tensors = ctx.last_step_outputs
do_validation = bool(validation_steps)
if model._compile_distribution:
distributed_training_utils._copy_weights_to_distributed_model(model, mode)
callbacks = cbks.configure_callbacks(
callbacks,
model,
do_validation=do_validation,
epochs=epochs,
steps_per_epoch=steps_per_epoch,
verbose=verbose,
count_mode='steps',
mode=mode)
# Calculate the steps each time on the device.
if use_steps:
steps_to_run = ([current_strategy.extended.steps_per_run] *
(steps_per_epoch //
current_strategy.extended.steps_per_run))
if steps_per_epoch % current_strategy.extended.steps_per_run:
steps_to_run.append(
steps_per_epoch % current_strategy.extended.steps_per_run)
target_steps = len(steps_to_run)
else:
target_steps = np.inf
callbacks._call_begin_hook(mode)
for epoch in range(initial_epoch, epochs):
distributed_training_utils._reset_metrics(model)
callbacks.on_epoch_begin(epoch)
epoch_logs = {}
step_index = 0
prev_step_count = None
current_step = 0
while current_step < target_steps:
step_count = steps_to_run[current_step] if use_steps else 1
batch_logs = {'batch': step_index, 'size': 1, 'num_steps': step_count}
callbacks._call_batch_hook(mode, 'begin', step_index, batch_logs)
if prev_step_count is None or step_count != prev_step_count:
steps_per_run.load(step_count, K.get_session())
prev_step_count = step_count
try:
_, outputs = K.batch_get_value([train_op, output_tensors])
except errors.OutOfRangeError:
if use_steps:
logging.warning('Your dataset iterator ran out of data; '
'interrupting training. Make sure that your dataset '
'can generate at least `steps_per_epoch * epochs` '
'batches (in this case, %d batches).' %
steps_per_epoch * epochs)
else:
target_steps = current_step
logging.info('Dataset iterator ran out of data. Inferring the '
'value of `steps_per_epoch` as %s .' % target_steps)
distributed_training_utils.initialize_iterator(iterator,
current_strategy)
break
batch_logs.update(outputs)
callbacks._call_batch_hook(mode, 'end', step_index, batch_logs)
step_index = step_index + step_count
current_step += 1
if callbacks.model.stop_training:
break
if (do_validation and
training_utils.should_run_validation(validation_freq, epoch)):
logging.info('Running validation at fit epoch: %s', epoch)
if model._compile_distribution:
# Since we create a new clone from the original model we need to copy
# the weights back to the original model before we can run validation.
distributed_training_utils._copy_weights_to_original_model(
model, ModeKeys.TRAIN)
val_outs = experimental_tpu_test_loop( # pylint: disable=undefined-variable
model,
val_dataset,
steps=validation_steps,
verbose=verbose,
callbacks=callbacks)
if not isinstance(val_outs, list):
val_outs = [val_outs]
# Same labels assumed.
for label, val_out in zip(out_labels, val_outs):
epoch_logs['val_' + label] = val_out
callbacks.on_epoch_end(epoch, epoch_logs)
if callbacks.model.stop_training:
break
callbacks._call_end_hook(mode)
if model._compile_distribution:
# Copy the weights back from the replicated model to the original model.
distributed_training_utils._copy_weights_to_original_model(
model, ModeKeys.TRAIN)
scope.__exit__(None, None, None)
return model.history
def experimental_tpu_predict_loop(model,
dataset,
verbose=0,
steps=None,
callbacks=None):
"""Predict loop for predicting with TPU tf.distribute.Strategy.
Arguments:
model: Keras Model instance.
dataset: Dataset for input data.
verbose: Integer, Verbosity mode 0 or 1.
steps: Total number of steps (batches of samples)
before declaring `_predict_loop` finished.
Ignored with the default value of `None`.
callbacks: List of callbacks to be called during training
Returns:
Array of predictions (if the model has a single output)
or list of arrays of predictions
(if the model has multiple outputs).
"""
mode = ModeKeys.PREDICT
dataset_fully_shaped = dist_utils.is_dataset_shape_fully_defined(dataset)
padding_handler = None
if not dataset_fully_shaped:
# TODO(hongjunchoi): Investigate whether operations from
# PartialBatchPaddingHandler are unnecessarily pruned out
# during graph optimization.
padding_handler = padding_util.PartialBatchPaddingHandler(
model._feed_output_shapes)
batch_size, _, prefetch_buffer = input_lib._get_dataset_attributes(
dataset)
padding_handler.padded_batch_size = batch_size
padding_handler.padding_mask = dataset.reduce(
padding_handler.padding_mask, padding_handler.update_mask)
dataset = dataset.map(padding_handler.pad_batch)
dataset = dataset.unbatch()
# Upon this point, it is guaranteed that the dataset does not
# have partial batches. Thus, we set `drop_remainder=True` to
# get static shape information about the elements in the dataset.
dataset = dataset.batch(batch_size, drop_remainder=True)
if prefetch_buffer is not None:
dataset = dataset.prefetch(prefetch_buffer)
current_strategy = model._distribution_strategy
iterator = dist_utils.get_iterator(dataset, current_strategy)
scope = dist_utils.distributed_scope(strategy=current_strategy,
learning_phase=0)
scope.__enter__()
def _predict_step_fn(inputs):
"""A fn that returns output of single prediction step."""
(distribution_strategy_context.get_replica_context().merge_call(
_build_model, args=(model, mode, inputs)))
(_, outputs, updates, _) = _per_replica_execution_function(
dist_utils.get_distributed_model(model, mode), mode)
with ops.control_dependencies([updates]):
return [array_ops.identity(out) for out in outputs]
# TODO(hongjunchoi): When numpy array is passed as an input to `predict()`
# use numpy arrays directly to avoid cumulating unnecessary input pipeline
# ops.
predict_input_data = iterator.get_next()
per_replica_outputs = current_strategy.run(_predict_step_fn,
args=(predict_input_data, ))
output_tensors = dist_utils.flatten_per_replica_values(
current_strategy, per_replica_outputs)
if verbose >= 1:
progbar = Progbar(target=steps)
if model._compile_distribution:
dist_utils._copy_weights_to_distributed_model(model, mode)
dist_utils._reset_metrics(model)
callbacks = cbks.configure_callbacks(callbacks,
model,
do_validation=False,
epochs=1,
steps_per_epoch=steps,
verbose=verbose,
count_mode='steps',
mode=mode)
callbacks._call_begin_hook(mode)
# Since we do not know how many samples we will see, we cannot pre-allocate
# the returned Numpy arrays. Instead, we store one array per batch seen
# and concatenate them upon returning.
num_model_outputs = len(model.output_names)
unconcatenated_outs = [[] for _ in range(num_model_outputs)]
if steps is not None:
target_steps = steps
else:
raise ValueError(
'Number of steps could not be inferred from the data, '
'please pass the steps argument.')
current_step = 0
while current_step < target_steps:
batch_logs = {'batch': current_step, 'size': 1}
callbacks._call_batch_hook(mode, 'begin', current_step, batch_logs)
try:
predict_ops = control_flow_ops.group(output_tensors)
_, batch_outs = K.batch_get_value([predict_ops, output_tensors])
except errors.OutOfRangeError:
warning_msg = (
'Make sure that your dataset can generate at least '
'`steps` batches (in this case, {} batches).'.format(steps))
logging.warning('Your dataset iterator ran out of data; '
'interrupting evaluation. ' + warning_msg)
break
# TODO(priyag): maybe need to unwrap the outputs first for MirroredStrategy.
for i in range(num_model_outputs):
output_start_index = i * current_strategy.num_replicas_in_sync
output_end_index = (output_start_index +
current_strategy.num_replicas_in_sync)
single_model_output = batch_outs[
output_start_index:output_end_index]
unconcatenated_outs[i].extend(single_model_output)
batch_logs = cbks.make_logs(model, batch_logs, batch_outs, mode)
callbacks._call_batch_hook(mode, 'end', current_step, batch_logs)
if verbose == 1:
progbar.update(current_step + 1)
current_step += 1
if verbose >= 1:
# Progress bar finishes at the end.
progbar.update(current_step)
callbacks._call_end_hook(mode)
scope.__exit__(None, None, None)
if len(unconcatenated_outs) == 1:
prediction_result = np.concatenate(unconcatenated_outs[0], axis=0)
else:
prediction_result = [
np.concatenate(out, axis=0) for out in unconcatenated_outs
]
if padding_handler:
prediction_result = padding_handler.apply_mask(prediction_result)
return prediction_result
def fit(self,
x=None,
y=None,
batch_size=None,
epochs=1,
verbose=1,
auto_switch=True,
retry_fit=True,
absorb=True,
train_after_switch=True,
callbacks=None,
validation_split=0.,
validation_data=None,
shuffle=True,
class_weight=None,
sample_weight=None,
initial_epoch=0,
steps_per_epoch=None,
validation_steps=None,
validation_batch_size=None,
validation_freq=1,
max_queue_size=10,
workers=1,
use_multiprocessing=False,
revert_after_fit=False):
"""
Custom fit function for the context model
auto_switch: Enable/disable autonomous context switching
train_after_switch:
retry_fit: Locate the next fitting context by re-performing fit.
absorb: Reset the switch sequence counter upon successful training.
This is mainly used to maintain switch sequencing for temporally-extended tasks
revert_after_fit This is a debug parameter to revert weights after performing a fit. This is used
to calculate the context deltas without incorrectly learning while auto switching
is disabled
"""
training._keras_api_gauge.get_cell('fit').set(True)
# Legacy graph support is contained in `training_v1.Model`.
version_utils.disallow_legacy_graph('Model', 'fit')
self._assert_compile_was_called()
self._check_call_args('fit')
if validation_split:
# Create the validation data using the training data. Only supported for
# `Tensor` and `NumPy` input.
(x, y, sample_weight), validation_data = (
data_adapter.train_validation_split(
(x, y, sample_weight),
validation_split=validation_split,
shuffle=False))
with self.distribute_strategy.scope(
), training_utils.RespectCompiledTrainableState(self):
# Creates a `tf.data.Dataset` and handles batch and epoch iteration.
data_handler = WindowedDataHandler(
x=x,
y=y,
sample_weight=sample_weight,
batch_size=batch_size,
steps_per_epoch=steps_per_epoch,
initial_epoch=initial_epoch,
epochs=epochs,
shuffle=shuffle,
class_weight=class_weight,
max_queue_size=max_queue_size,
workers=workers,
use_multiprocessing=use_multiprocessing,
model=self)
# Container that configures and calls `tf.keras.Callback`s.
if not isinstance(callbacks, callbacks_module.CallbackList):
callbacks = callbacks_module.CallbackList(
callbacks,
add_history=True,
add_progbar=bool(verbose & Verbosity.Progress),
model=self,
verbose=verbose,
epochs=epochs,
steps=data_handler.inferred_steps)
self.stop_training = False
train_function = self.make_train_function()
callbacks.on_train_begin()
self.initialize_fit()
# Handle fault-tolerance for multi-worker.
# TODO(omalleyt): Fix the ordering issues that mean this has to
# happen after `callbacks.on_train_begin`.
data_handler._initial_epoch = (
self._maybe_load_initial_epoch_from_ckpt(initial_epoch))
for epoch, window_iterator in data_handler.enumerate_epochs():
self.reset_metrics()
callbacks.on_epoch_begin(epoch)
dataset = tf.data.Dataset.zip(next(window_iterator))
switched_during_epoch = False # Indicate if the model has attempted at least one switch during this epoch
switched = True # Indicate if the model switched on the most recent fit iteration
weights = backend.batch_get_value(self.trainable_variables)
# Perform a 'fit call'. Assuming retry_fit, this call is re-attempted after each switch until a context fits
while switched and (retry_fit or not switched_during_epoch):
self.initialize_epoch(epoch)
iterator = iter(dataset)
# Perform a fit call
with data_handler.catch_stop_iteration():
for step in data_handler.steps():
with traceme.TraceMe('TraceContext',
graph_type='train',
epoch_num=epoch,
step_num=step,
batch_size=batch_size):
callbacks.on_train_batch_begin(step)
tmp_logs = train_function(iterator)
# Catch OutOfRangeError for Datasets of unknown size.
# This blocks until the batch has finished executing.
# TODO(b/150292341): Allow multiple async steps here.
if not data_handler.inferred_steps:
context.async_wait()
logs = tmp_logs # No error, now safe to assign to logs.
callbacks.on_train_batch_end(step, logs)
switched = not self.update_and_switch(
epoch, auto_switch, absorb, retry_fit, verbose)
switched_during_epoch |= switched
# If a switch occurred, we need to restore the weights
if switched or (switched_during_epoch
and not train_after_switch
) or revert_after_fit:
backend.batch_set_value(
zip(self.trainable_variables, weights))
self.reset_metrics()
epoch_logs = copy.copy(logs)
# Run validation.
if validation_data and self._should_eval(
epoch, validation_freq):
val_x, val_y, val_sample_weight = (
data_adapter.unpack_x_y_sample_weight(validation_data))
val_logs = self.evaluate(
x=val_x,
y=val_y,
sample_weight=val_sample_weight,
batch_size=validation_batch_size or batch_size,
steps=validation_steps,
callbacks=callbacks,
max_queue_size=max_queue_size,
workers=workers,
use_multiprocessing=use_multiprocessing,
return_dict=True)
val_logs = {
'val_' + name: val
for name, val in val_logs.items()
}
epoch_logs.update(val_logs)
callbacks.on_epoch_end(epoch, epoch_logs)
if self.stop_training:
break
callbacks.on_train_end()
return self.history
def save_model(model, filepath, overwrite=True, include_optimizer=True):
"""Saves a model to a HDF5 file.
The saved model contains:
- the model's configuration (topology)
- the model's weights
- the model's optimizer's state (if any)
Thus the saved model can be reinstantiated in
the exact same state, without any of the code
used for model definition or training.
Arguments:
model: Keras model instance to be saved.
filepath: One of the following:
- String, path where to save the model
- `h5py.File` object where to save the model
overwrite: Whether we should overwrite any existing
model at the target location, or instead
ask the user with a manual prompt.
include_optimizer: If True, save optimizer's state together.
Raises:
ImportError: if h5py is not available.
"""
if h5py is None:
raise ImportError('`save_model` requires h5py.')
from tensorflow.python.keras import __version__ as keras_version # pylint: disable=g-import-not-at-top
if not isinstance(filepath, h5py.File):
# If file exists and should not be overwritten.
if not overwrite and os.path.isfile(filepath):
proceed = ask_to_proceed_with_overwrite(filepath)
if not proceed:
return
f = h5py.File(filepath, mode='w')
opened_new_file = True
else:
f = filepath
opened_new_file = False
try:
f.attrs['keras_version'] = str(keras_version).encode('utf8')
f.attrs['backend'] = K.backend().encode('utf8')
f.attrs['model_config'] = json.dumps(
{
'class_name': model.__class__.__name__,
'config': model.get_config()
},
default=serialization.get_json_type).encode('utf8')
model_weights_group = f.create_group('model_weights')
model_layers = model.layers
save_weights_to_hdf5_group(model_weights_group, model_layers)
if include_optimizer and hasattr(model, 'optimizer'):
if isinstance(model.optimizer, optimizers.TFOptimizer):
logging.warning(
'TensorFlow optimizers do not '
'make it possible to access '
'optimizer attributes or optimizer state '
'after instantiation. '
'As a result, we cannot save the optimizer '
'as part of the model save file.'
'You will have to compile your model again after loading it. '
'Prefer using a Keras optimizer instead '
'(see keras.io/optimizers).')
else:
f.attrs['training_config'] = json.dumps(
{
'optimizer_config': {
'class_name': model.optimizer.__class__.__name__,
'config': model.optimizer.get_config()
},
'loss': model.loss,
'metrics': model.metrics,
'sample_weight_mode': model.sample_weight_mode,
'loss_weights': model.loss_weights,
},
default=serialization.get_json_type).encode('utf8')
# Save optimizer weights.
symbolic_weights = getattr(model.optimizer, 'weights')
if symbolic_weights:
optimizer_weights_group = f.create_group(
'optimizer_weights')
weight_values = K.batch_get_value(symbolic_weights)
weight_names = []
for w, val in zip(symbolic_weights, weight_values):
name = str(w.name)
weight_names.append(name.encode('utf8'))
optimizer_weights_group.attrs[
'weight_names'] = weight_names
for name, val in zip(weight_names, weight_values):
param_dset = optimizer_weights_group.create_dataset(
name, val.shape, dtype=val.dtype)
if not val.shape:
# scalar
param_dset[()] = val
else:
param_dset[:] = val
f.flush()
finally:
if opened_new_file:
f.close()
