title: Keras - no good way to stop and resume training?
body: <p>After a lot of research, it seems like there is no good way to <em>properly</em> stop and resume training using a Tensorflow 2 / Keras model. This is true whether you are using <code>model.fit()</code> <em>or</em> using a custom training loop.</p>
  <p>There seem to be 2 supported ways to save a model while training:</p>
  <ol>
  <li><p>Save just the weights of the model, using <code>model.save_weights()</code> or <code>save_weights_only=True</code> with <code>tf.keras.callbacks.ModelCheckpoint</code>. This seems to be preferred by most of the examples I've seen, however it has a number of major issues:</p>
  <ul>
  <li>The optimizer state is not saved, meaning training resumption will not be correct.</li>
  <li>Learning rate schedule is reset - this can be catastrophic for some models.</li>
  <li>Tensorboard logs go back to step 0 - making logging essentually useless unless complex workarounds are implemented.</li>
  </ul>
  </li>
  <li><p>Save the entire model, optimizer, etc. using <code>model.save()</code> or <code>save_weights_only=False</code>. The optimizer state is saved (good) but the following issues remain:</p>
  <ul>
  <li>Tensorboard logs still go back to step 0</li>
  <li>Learning rate schedule is still reset (!!!)</li>
  <li>It is impossible to use custom metrics.</li>
  <li>This doesn't work at all when using a custom training loop - custom training loops use a non-compiled model, and saving/loading a non-compiled model doesn't seem to be supported.</li>
  </ul>
  </li>
  </ol>
  <p>The best workaround I've found is to use a custom training loop, manually saving the step. This fixes the tensorboard logging, and the learning rate schedule can be fixed by doing something like <code>keras.backend.set_value(model.optimizer.iterations, step)</code>. However, since a full model save is off the table, the optimizer state is not preserved. I can see no way to save the state of the optimizer independently, at least without a lot of work. And messing with the LR schedule as I've done feels messy as well.</p>
  <p>Am I missing something? How are people out there saving/resuming using this API?</p>