tensorflow: Restoring Keras model fails inside a distribution strategy scope

I think the optimizer weights issue is because sometimes optimizer weights are only created during the first step. So when you try to load the weights before training starts, weights have not yet been created. One workaround can be to just do 1 step first and then load the weights. @k-w-w do you know what could be a better alternative for loading optimizer weights ?

Thanks, it seems to work with the 2.0 preview version.

For TF 1.14, I’ve found a workaround, which I’ll document here for anyone else coming across this:

1. Save the model with tf.keras.models.save_model.
2. Load a copy of the model outside the strategy scope using tf.keras.models.load_model. Retrieve its weights using model.get_weights() and model.optimizer.get_weights().
3. Create and compile a new copy of the model inside the scope, in the same way as you did the first time.
4. Now, in a callback, copy the loaded weights to the newly compiled model and its optimizer. The callback is key: the optimizer weights are not created before the call to model.fit, so the copying cannot be done in the normal flow of code, but the on_train_begin callback runs after the weights are created and before any training has happened.

An example of what the callback might look like:

class LoadWeightsCallback(tf.keras.callbacks.Callback):
    _chief_worker_only = False

    def __init__(self, weights, optimizer_weights):
        self.weights = weights
        self.optimizer_weights = optimizer_weights

    def on_train_begin(self, logs=None):
        self.model.set_weights(self.weights)
        self.model.optimizer.set_weights(self.optimizer_weights)

For TF 2.0, I have tf.version.VERSION == '2.0.0-beta1' and tf.version.GIT_VERSION == 'v2.0.0-beta0-17-g8e423e3'. I installed via pip install tensorflow-gpu==2.0.0b1. That is the latest released version on both GitHub and PyPI, as far as I can tell.

If I understand TF correctly, I don’t need to store the graph, as I am restoring in Python with access to the same code that created the model originally. So perhaps save_model and load_model are overkill. However, I do need the state of the optimizer, which rules out using just model.{save,load}_weights. I also attempted to load one copy of the model outside the scope, construct another copy inside the scope, and then use model.{get,set}_weights and model.optimizer.{get,set}_weights, but it seems that the optimizer weights are not created until model.fit has been called with data, so that runs into issues around mismatch of expected number of weights.

In TF 1.14, I also tried tf.contrib.saved_model.{save,load}_keras_model, which fails in a different way:

import numpy as np, tensorflow as tf

strategy = tf.distribute.MirroredStrategy()
path = "/tmp/model"

with strategy.scope():
    model = tf.keras.models.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
    model.compile(optimizer=tf.keras.optimizers.SGD(), loss=tf.keras.metrics.mse)
    model.fit(np.array([[1]]), np.array([[1]]))
    tf.contrib.saved_model.save_keras_model(model, path)

    model = tf.contrib.saved_model.load_keras_model(path)
    model.compile(optimizer=tf.keras.optimizers.SGD(), loss=tf.keras.metrics.mse)
    model.fit(np.array([[1]]), np.array([[1]]))

gives this exception inside the second model.fit:

  File ".../tensorflow/python/keras/engine/training.py", line 649, in fit
    validation_freq=validation_freq)
  File ".../tensorflow/python/keras/engine/training_distributed.py", line 143, in fit_distributed
    steps_name='steps_per_epoch')
  File ".../tensorflow/python/keras/engine/training_arrays.py", line 274, in model_iteration
    batch_outs = f(actual_inputs)
  File ".../tensorflow/python/keras/backend.py", line 3292, in __call__
    run_metadata=self.run_metadata)
  File ".../tensorflow/python/client/session.py", line 1458, in __call__
    run_metadata_ptr)
tensorflow.python.framework.errors_impl.FailedPreconditionError: Error while reading resource variable dense_2/kernel from Container: localhost. This could mean that the variable was uninitialized. Not found: Resource localhost/dense_2/kernel/N10tensorflow3VarE does not exist.
         [[{{node dense_3/MatMul/ReadVariableOp}}]]

Hi - we don’t support saving with hdf5 format. However, you can save and restore with the standard TF format - just remove the hdf5 extension from the file path. See https://www.tensorflow.org/beta/tutorials/distribute/save_and_load for more information.