TensorFlow
Object Detection
Machine Learning
Custom Dataset
Deep Learning

Train Tensorflow Object Detection on own dataset

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

Training TensorFlow object detection on your own dataset is mostly a data and pipeline problem, not just a model-selection problem. The model can only learn what your annotations, label map, and input pipeline make possible. A successful setup usually comes down to four things: clean labels, the right record format, a correct pipeline config, and patience with debugging before long training runs.

Prepare And Label The Dataset

The first requirement is a labeled dataset where every image has bounding boxes and class names. The labels must be consistent.

For example, if one annotator uses car and another uses cars, you have created two classes whether you meant to or not. Annotation quality usually matters more than choosing between two similar pre-trained architectures.

A common split is:

  • training images,
  • validation images,
  • optional held-out test images.

The validation set should reflect the same real-world distribution you care about at inference time.

Create A Label Map

TensorFlow’s Object Detection API uses a label map to assign integer IDs to class names.

protobuf
1item {
2  id: 1
3  name: "cat"
4}
5
6item {
7  id: 2
8  name: "dog"
9}

Keep IDs stable. If you later change class numbering carelessly, checkpoints and exported models become harder to reason about.

Convert The Dataset To TFRecord

Most training pipelines use TFRecord files for performance and pipeline consistency. The exact converter depends on your annotation format, but the end goal is usually a train.record and val.record pair.

Once those files exist, training becomes much easier to automate and debug than reading many small annotation files directly.

Start From A Pretrained Model

Training from scratch is possible, but transfer learning is usually the right default. Pick a pre-trained detector close to your accuracy-speed tradeoff and fine-tune it.

In the pipeline config, the key fields you usually customize are:

  • 'num_classes,'
  • 'fine_tune_checkpoint,'
  • training record path,
  • validation record path,
  • label map path,
  • batch size,
  • total training steps.

A small config excerpt often looks like this:

protobuf
1model {
2  faster_rcnn {
3    num_classes: 2
4  }
5}
6
7train_config {
8  batch_size: 4
9  fine_tune_checkpoint: "pretrained_model/checkpoint/ckpt-0"
10}

The exact model block varies by architecture, but the same idea holds.

Launch Training Carefully

A typical command looks like this:

bash
1python model_main_tf2.py \
2  --pipeline_config_path=training/pipeline.config \
3  --model_dir=training/output \
4  --alsologtostderr

Before you start a long run, verify:

  • the dataset paths exist,
  • 'num_classes matches the label map,'
  • the checkpoint path is correct,
  • one batch can be read without errors.

This saves a lot of wasted training time.

Watch Evaluation Early

Do not wait until the end of training to discover that labels were wrong or boxes were malformed. Run evaluation against the validation split early and often.

Useful signals include:

  • loss trends,
  • mAP trends,
  • visualized predictions,
  • class-wise failure patterns.

A model can show decreasing loss while still being practically unusable if the annotations or class definitions are inconsistent.

Data Problems Usually Beat Model Problems

When custom training disappoints, the cause is often one of these:

  • too little data,
  • inconsistent annotation boxes,
  • class imbalance,
  • domain mismatch between training images and real deployment images,
  • pipeline config not aligned with the dataset.

It is common to spend time switching model architectures when the real issue is annotation quality.

Export The Trained Model

Once you are satisfied with validation quality, export the trained checkpoint for inference.

bash
1python exporter_main_v2.py \
2  --input_type image_tensor \
3  --pipeline_config_path training/pipeline.config \
4  --trained_checkpoint_dir training/output \
5  --output_directory exported-model

That produces a SavedModel-style export suitable for later inference workflows.

Common Pitfalls

  • Inconsistent class names across annotations.
  • Mismatch between num_classes and the label map.
  • Bad checkpoint path or wrong pre-trained checkpoint type.
  • Training for hours before verifying that the pipeline reads records correctly.
  • Blaming the model choice when the real issue is annotation quality or dataset size.

Summary

  • Successful TensorFlow object detection training starts with clean labeled data.
  • Build a correct label map and TFRecord pipeline before long training runs.
  • Fine-tuning a pre-trained detector is usually better than training from scratch.
  • Validate early with both metrics and visual inspection.
  • Most failures come from dataset and config mistakes more often than from the architecture itself.

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.