def save(
self,
model: tf.keras.models.Model,
save_dir: str,
save_sub_format: SaveSubFormat = SaveSubFormat.TF,
) -> None:
"""
Save the model using the correct format and sub-format.
Args:
model (:py:class:`tf.keras.models.Model`): model to Save.
save_dir (str): path of the file in which to save the model.
save_sub_format (:py:class:`ashpy.callbacks.save_callback.SaveSubFormat`): sub-format
of the save operation.
"""
if SaveFormat.WEIGHTS & self:
save_dir = self._initialize_dirs(save_dir, SaveFormat.WEIGHTS,
save_sub_format)
model.save_weights(save_dir, save_format=save_sub_format.value)
if SaveFormat.MODEL & self:
save_dir = self._initialize_dirs(save_dir, SaveFormat.MODEL,
save_sub_format)
model.save(save_dir, save_format=save_sub_format.value)
if not (SaveFormat.MODEL & self) | (SaveFormat.WEIGHTS & self):
raise NotImplementedError(
"No implementation of `save` method for the current SaveFormat"
)
def train_model(model: tf.keras.models.Model,
train_gen: DataGenerator,
validation_gen: DataGenerator,
epochs=10,
steps=4000,
checkpoints_path="checkpoints"):
mk_dir(checkpoints_path)
model_dir = os.path.join(checkpoints_path, tb.time_stamp)
tf.saved_model.save(model, model_dir + "/model")
optimizer = tf.keras.optimizers.Adadelta()
global_steps = tf.Variable(0, trainable=False, dtype=tf.int64)
validation_steps = tf.Variable(0, trainable=False, dtype=tf.int64)
tb.generate_model_graph(model, train_gen.get_sample_data())
tb.launch_tensorboard()
for epoch in range(epochs):
for step in range(steps):
global_steps.assign_add(1)
image_data, target_y_e, target_y_g, target_y_a = next(
train_gen.get_data())
with tf.GradientTape() as tape:
pred_y_e, pred_y_g, pred_y_a = model(image_data)
l_e, l_g, l_a = compute_loss(pred_y_e, pred_y_g, pred_y_a,
target_y_e, target_y_g,
target_y_a)
total_loss = l_e + l_g + l_a
gradients = tape.gradient(total_loss,
model.trainable_variables)
optimizer.apply_gradients(
zip(gradients, model.trainable_variables))
print("=> epoch %d step %d train_loss: %.6f" %
(epoch + 1, step + 1, total_loss.numpy()))
tb.add_scalar("train/total_loss",
total_loss,
step=global_steps)
tb.add_scalar("train/emotion_loss", l_e, step=global_steps)
tb.add_scalar("train/gender_loss", l_g, step=global_steps)
tb.add_scalar("train/age_loss", l_a, step=global_steps)
# validation step
if step % 500 == 0:
validation_steps.assign_add(1)
image_data, target_y_e, target_y_g, target_y_a = next(
validation_gen.get_data())
pred_y_e, pred_y_g, pred_y_a = model(image_data)
l_e, l_g, l_a = compute_loss(pred_y_e, pred_y_g, pred_y_a,
target_y_e, target_y_g,
target_y_a)
total_valid_loss = l_e + l_g + l_a
tb.add_scalar("valid_loss",
total_valid_loss,
step=validation_steps)
p_loss = int(round(total_loss.numpy(), 2) * 100)
model.save(f"{model_dir}/EGA_epoch_{epoch}_score_{p_loss}.model")
model.save_weights(f"{model_dir}/EGA_epoch_{epoch}_score_{p_loss}.h5")
def save_tensorflow_model(model: tf.keras.models.Model, path: str):
assert path != ""
os.makedirs(path, exist_ok=True)
model.save_weights(os.path.join(path, "parameters.h5"))
with open(os.path.join(path, "topology.json"), 'w') as f:
f.write(model.to_json())
print(f"Saved model to {path}")
def save_keras_model(model: tf.keras.models.Model, json_file: str,
weights_file: str) -> None:
"""Save keras model's architecture and weights to disk
Args:
model: Instance of tf.keras model.
json_file: str, Path to save model's architecture info.
weights_file: str. Path to save model's weights.
Returns:
"""
model_json = model.to_json()
with open(json_file, 'w') as writer:
writer.write(model_json)
model.save_weights(weights_file)
logging.info('Saved model to disk')
def save(
self,
model: tf.keras.models.Model,
save_dir: Path,
save_sub_format: SaveSubFormat = SaveSubFormat.TF,
) -> None:
"""
Save the model using the correct format and sub-format.
Args:
model (:py:class:`tf.keras.models.Model`): model to Save.
save_dir (:class:`pathlib.Path`): path of the file in which to save the model.
save_sub_format (:py:class:`ashpy.callbacks.save_callback.SaveSubFormat`): sub-format
of the save operation.
"""
if SaveFormat.WEIGHTS & self:
save_dir = self._initialize_dirs(save_dir, SaveFormat.WEIGHTS,
save_sub_format)
# NOTE: Keras (TF 2.1.0) checks for h5 file using endswith attribute.
# Explicit conversion to strings is required
model.save_weights(str(save_dir),
save_format=save_sub_format.value)
if SaveFormat.MODEL & self:
save_dir = self._initialize_dirs(save_dir, SaveFormat.MODEL,
save_sub_format)
# NOTE: TensorFlow 2.1.0 wanth either binary or unicod string.
# Explicit conversion to strings is required
model.save(str(save_dir), save_format=save_sub_format.value)
if not (SaveFormat.MODEL & self) | (SaveFormat.WEIGHTS & self):
raise NotImplementedError(
"No implementation of `save` method for the current SaveFormat"
)
def stochastic_from_keras(
model: tf.keras.models.Model,
input_tensors=None,
clone_function=None,
expect_determinism=False,
temp_weights_path="tmp/weights",
):
"""
Creates a stochastic instance from a given `tf.keras.models.Sequential` model:
The new model will have the same structure (layers) and weights as the passed model.
All stochastic layers (e.g. tf.keras.layers.Dropout) will be used for randomization during randomized predictions.
If no stochastic layers are present, a ValueError is thrown.
The raising of the error can be suppressed by setting `expect_determinism` to true.
If your model contains custom layers, you can pass a function to `clone_function` to clone your custom layers,
or place the annotation `@tf.keras.utils.register_keras_serializable()` on your custom layers,
and make sure the `get_config` and `from_config` methods are implemented.
(uncertainty wizard will serialize and deserialize all layers).
:param model: The model to copy. Remains unchanged.
:param input_tensors: Optional tensors to use as input_tensors for new model. See the corresponding parameter in `tf.keras.models.clone_model` for details.
:param _clone_function: Optional function to use to clone layers. Will be applied to all layers except input layers and stochastic layers. See the corresponding parameter in `tf.keras.models.clone_model` for more details.
:param expect_determinism: If True, deterministic models (e.g. models without stochastic layers) are accepted and no ValueError is thrown.
:param temp_weights_path: The model weights are temporarily saved to the disk at this path. Folder is deleted after successful completion.
:return: A newly created stochastic model
"""
# _clone_function is some layer cloning behavior that can be specified by the user
# If none is specified, we use keras default (see `tf.keras.models.clone_model`)
if clone_function is None:
def _clone_function(layer):
return layer.__class__.from_config(layer.get_config())
# We wrap the users (or default) clone function in a clone function
# that replaces stochastic layers with uncertainty wizard stochastic layers
stochastic_mode = StochasticMode()
is_stochastic_layer = []
def _uncertainty_wizard_aware_clone_function(layer):
new_layer = Stochastic._replace_layer_if_possible(
layer, stochastic_mode=stochastic_mode
)
if new_layer == layer:
# Layer was not mapped to an uncertainty wizard layer, thus the default clone function is applied
new_layer = _clone_function(layer)
is_stochastic_layer.append(False)
else:
is_stochastic_layer.append(True)
return new_layer
# Clone the keras model to become the new inner model
new_inner = tf.keras.models.clone_model(
model=model,
input_tensors=input_tensors,
clone_function=_uncertainty_wizard_aware_clone_function,
)
new_inner.stochastic_mode_tensor = stochastic_mode.as_tensor()
if not expect_determinism and not any(is_stochastic_layer):
raise ValueError(
"The passed model had no stochastic layers."
"If that is intended (and you do not plan to use any sampling based quantifiers)"
"you can set the flag `expect_determinism = True`, i.e., "
"calling `SequentialStochastic.clone_from_keras(keras_model,expect_determinism = True)`"
)
# Restore the Weights
model.save_weights(temp_weights_path)
new_inner.load_weights(temp_weights_path)
# Remove temporarily stored weights
shutil.rmtree(temp_weights_path, ignore_errors=True)
# Put the wrapper around the new model
if isinstance(model, tf.keras.models.Sequential):
target_class = StochasticSequential
else:
target_class = StochasticFunctional
# Consenting Adults: The _wrap method is intended to be used here
# but declared private as it is not intended to be used by the uwiz user
return target_class._wrap(
inner=new_inner, stochastic_mode_tensor=stochastic_mode.as_tensor()
)