Is it possible to restore global_step in keras?

user_6720913 0
Viewed 530

keras can save model with tf.keras.models.save_model(), then restore model from tf.keras.models.load_model()

Here is a scenario,

  1. train a model x step. save the weights to checkpoint file. save_format='tf'
  2. restore weight from checkpoint file, train another x step.
  3. In tensorflow the starting global step of 2nd training is x+1 step, but in keras, it seems that it reset the global steps.

So far, this behavior cause me 2 problems.

  1. tensorboad is messing up, the 2 training run all start from step 0. tensorboard
  2. I need save some information into file during the training, the file name will depends on the train step, so I can save information for different steps, eg: weight_1000steps, weight_2000steps, etc. But since step will restart if I start a new training, so some old file will be overwritten.

How can I restore the training step/epoch from model file in keras? (tf.__version_ is '2.3.1')

tensorflow2.0 keras tensorflow
1 Answers

In Keras, the optimizers keep track of the iterations themselves and their iteration counter is persisted in the checkpoints. It is available as optimizer attribute model.optimizer.iterations.

You can compute the corresponding epoch number from this and use it to adjust the epoch number for tensorboard using the parameter init_epoch of model.fit().

I haven't found a documented way to adjust the step number for tensorboard, but at least for current tensorflow (2.4) it can be done by altering the model's (internal!) field _train_counter with a callback passed in the callbacks parameter of model.fit():

class UseOptimizerIterationAsTrainStep(tf.keras.callbacks.Callback):
    def set_model(self, model):
        self.model = model

    def on_train_begin(self, logs=None):
        self.model._train_counter.assign(self.model.optimizer.iterations)

Of course this may break again in future tensorflow versions.

user_97021420
Related