TensorFlow
Object Detection
Model Checkpoints
Machine Learning
API Usage

How to store best models checkpoints, not only newest 5, in Tensorflow Object Detection API?

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

When training object detection models, saving only the latest checkpoints is often not enough. The most recent checkpoint might not be the best one, especially if validation metrics fluctuate late in training. A better strategy is to keep both recent recovery checkpoints and top performing checkpoints ranked by evaluation metrics.

Why Latest Checkpoints and Best Checkpoints Solve Different Problems

Latest checkpoints help you resume interrupted training quickly. Best checkpoints help you deploy a model that actually performs well on validation data. These are separate goals and should be managed separately.

In TensorFlow workflows, teams usually need both:

  • rolling checkpoints for fault tolerance
  • metric ranked checkpoints for model selection

If you rely on only one policy, you usually lose either recoverability or quality.

Baseline Checkpoint Management in TensorFlow

tf.train.CheckpointManager provides rolling retention with max_to_keep. This is useful for saving the newest checkpoints during training.

python
1import tensorflow as tf
2from pathlib import Path
3
4model = tf.keras.Sequential([
5    tf.keras.layers.Input(shape=(32,)),
6    tf.keras.layers.Dense(64, activation='relu'),
7    tf.keras.layers.Dense(1)
8])
9
10optimizer = tf.keras.optimizers.Adam(1e-3)
11ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=optimizer, model=model)
12manager = tf.train.CheckpointManager(ckpt, directory='ckpts/latest', max_to_keep=5)
13
14# simulate periodic save
15for _ in range(7):
16    save_path = manager.save()
17    ckpt.step.assign_add(1)
18    print('saved:', save_path)

After seven saves, only the newest five remain in ckpts/latest.

Keep Top K Checkpoints by Validation Metric

To preserve best models, add a second retention flow keyed by validation score. The simple pattern below copies checkpoint files to a dedicated best folder and tracks ranking in JSON.

python
1import json
2import shutil
3from pathlib import Path
4
5class TopKCheckpointKeeper:
6    def __init__(self, best_dir: str, k: int = 3):
7        self.best_dir = Path(best_dir)
8        self.best_dir.mkdir(parents=True, exist_ok=True)
9        self.k = k
10        self.meta_file = self.best_dir / 'scores.json'
11        self.records = self._load()
12
13    def _load(self):
14        if self.meta_file.exists():
15            return json.loads(self.meta_file.read_text())
16        return []
17
18    def _save(self):
19        self.meta_file.write_text(json.dumps(self.records, indent=2))
20
21    def consider(self, score: float, ckpt_prefix: str):
22        candidate_name = Path(ckpt_prefix).name
23        target_prefix = self.best_dir / candidate_name
24
25        for suffix in ['.index', '.data-00000-of-00001']:
26            shutil.copyfile(ckpt_prefix + suffix, str(target_prefix) + suffix)
27
28        self.records.append({'score': score, 'prefix': candidate_name})
29        self.records.sort(key=lambda x: x['score'], reverse=True)
30        self.records = self.records[:self.k]
31
32        keep = {r['prefix'] for r in self.records}
33        for p in self.best_dir.glob('ckpt-*'):
34            prefix = p.name.split('.data')[0].split('.index')[0]
35            if prefix not in keep:
36                if p.exists():
37                    p.unlink()
38
39        self._save()

Use this keeper right after each evaluation pass.

python
1val_map = 0.412  # example metric from evaluator
2latest_prefix = manager.latest_checkpoint
3keeper = TopKCheckpointKeeper('ckpts/best', k=3)
4keeper.consider(val_map, latest_prefix)

Integrating with TensorFlow Object Detection API

Object Detection API training often uses separate training and evaluation binaries. A practical setup is:

  • training job writes rolling checkpoints in one directory
  • evaluation job computes mAP for each checkpoint
  • evaluation callback copies top checkpoints to best directory

If you use a custom training script, you can run evaluation every N steps and call consider directly. If you use pipeline based scripts, parse evaluator output and trigger retention logic as a post evaluation step.

Keep the metric name stable. If you switch between [email protected] and [email protected]:0.95 during one run, ranking becomes inconsistent.

Deployment Workflow with Best Checkpoints

At deployment time, load best checkpoint metadata and choose top ranked candidate. Optionally keep top two for A B comparisons.

python
1import json
2from pathlib import Path
3
4scores_path = Path('ckpts/best/scores.json')
5records = json.loads(scores_path.read_text())
6best = records[0]
7print('best score:', best['score'])
8print('best checkpoint prefix:', best['prefix'])

Then restore with your checkpoint object using the selected prefix.

Common Pitfalls

A common mistake is overwriting the same best checkpoint name repeatedly. Use unique checkpoint prefixes or include global step in file names so ranking logic can distinguish candidates.

Another issue is comparing metrics from different validation datasets. If dataset or preprocessing changes, old and new scores are not comparable, so top K selection becomes misleading.

A third issue is keeping only best checkpoints and deleting all recent checkpoints. If training crashes, you may lose resume capability. Keep a rolling latest directory and a separate best directory.

Teams also forget to store metadata about metric type and evaluation config. Without that context, future readers cannot trust score comparisons.

Summary

  • Keep latest and best checkpoints separately because they serve different goals
  • Use CheckpointManager for rolling retention and custom ranking for best models
  • Rank checkpoints with one stable validation metric
  • Preserve metadata so checkpoint scores are interpretable later
  • Maintain both recoverability and model quality in the same training pipeline

Related reading
Course
Beginner
27 lessons
10 hours
System Design Fundamentals

Build a strong foundation in designing scalable, reliable distributed systems.

View the 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