tensorflow
keras
tensorboard
resume training
machine learning

Resume Training tf.keras Tensorboard

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Resuming a tf.keras training run has two separate concerns: restoring model state and keeping your training logs coherent in TensorBoard. Restoring the weights alone lets you continue learning, but a true resume often also needs optimizer state, the correct starting epoch, and a logging setup that does not make the new run look unrelated.

The safest pattern is to save checkpoints during training, reload the model or weights after interruption, and call fit again with initial_epoch set to the last completed epoch.

Save Checkpoints During Training

If you want to resume later, training must periodically save recoverable state. A typical setup combines ModelCheckpoint with TensorBoard.

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Dense(32, activation="relu"),
5    tf.keras.layers.Dense(1)
6])
7
8model.compile(optimizer="adam", loss="mse")
9
10callbacks = [
11    tf.keras.callbacks.ModelCheckpoint(
12        filepath="checkpoints/epoch-{epoch:02d}.weights.h5",
13        save_weights_only=True
14    ),
15    tf.keras.callbacks.TensorBoard(log_dir="logs/run-1")
16]
17
18history = model.fit(
19    x_train,
20    y_train,
21    epochs=10,
22    callbacks=callbacks
23)

This produces weight checkpoints and TensorBoard event files. If training stops, you have something to restore from.

Resume With Weights And initial_epoch

When you resume from weights only, rebuild the model with the same architecture, load the saved weights, and continue training from the correct epoch count.

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Dense(32, activation="relu"),
5    tf.keras.layers.Dense(1)
6])
7
8model.compile(optimizer="adam", loss="mse")
9model.load_weights("checkpoints/epoch-10.weights.h5")
10
11callbacks = [
12    tf.keras.callbacks.TensorBoard(log_dir="logs/run-1")
13]
14
15model.fit(
16    x_train,
17    y_train,
18    initial_epoch=10,
19    epochs=20,
20    callbacks=callbacks
21)

initial_epoch=10 tells Keras that epochs 0 through 9 already happened. Without that, the resumed run starts counting again from zero and your logs become harder to interpret.

Weights-Only Resume Versus Full-State Resume

Weights-only resume is often enough, but it does not necessarily restore the optimizer's internal state exactly. For adaptive optimizers such as Adam, that means the resumed training may not be identical to an uninterrupted run.

If you want a closer continuation, save and reload the full model or use backup callbacks that preserve optimizer state:

python
1model.save("saved_model.keras")
2
3restored = tf.keras.models.load_model("saved_model.keras")
4restored.fit(
5    x_train,
6    y_train,
7    initial_epoch=10,
8    epochs=20,
9    callbacks=[tf.keras.callbacks.TensorBoard(log_dir="logs/run-1")]
10)

That is usually the better option when exact continuation matters more than lightweight checkpoint files.

Keep TensorBoard Logs Sensible

TensorBoard does not magically know two training sessions are conceptually one run. The log directory strategy is what determines how the graphs appear. If you want continuity, write the resumed training logs into the same run directory and keep initial_epoch aligned with the last completed epoch. If you want to compare before-and-after behavior as separate experiments, use a new log directory instead.

The key is to decide intentionally. Reusing the same directory with inconsistent epoch numbering creates the most confusing result because the charts no longer reflect a clean timeline.

Common Pitfalls

The most common mistake is restoring only the weights and expecting an exact continuation when the optimizer state was not restored. Another is forgetting initial_epoch, which makes the resumed run restart its counters and muddles TensorBoard charts. Developers also sometimes change the model architecture or compile settings before loading checkpoints, which can make restore fail or silently alter the training dynamics. Finally, be careful with log directories. If you reuse a directory accidentally for a different experiment, TensorBoard can merge unrelated events and make the run history misleading.

Summary

  • Save checkpoints during training if you want to resume later.
  • Reload weights or the full model before calling fit again.
  • Set initial_epoch to the last completed epoch so training history stays coherent.
  • Use full-model restore when optimizer state matters.
  • Reuse or separate TensorBoard log directories intentionally depending on whether you want continuity or comparison.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.