Training MobileNetV2-SSD: Target Assignment, Loss Functions & Training Tricks

In my other post, I covered the MobileNetV2-SSD architecture and the decisions that I made to make it deployable on edge devices. This post focuses on the training pipeline a bit more in-depth, including but not limited to target assignment, loss functions, hard negative mining etc.

Training Pipeline Overview

The full training loop involves multiple subsystems that operate in sequence on each batch.

Target Assignment (IoU Matching)

Before computing any loss, each of the model's ~8,732 prior anchor boxes must be assigned a label, either a ground-truth object class or "background." This is done through IoU-based bipartite matching.

For each prior and each ground-truth box in the image, we compute the Intersection over Union:

A prior is marked positive (matched to a ground-truth box) if its IoU exceeds 0.5, and negative (background) if its IoU falls below 0.4. Priors with IoU in the range [0.4, 0.5] are ignored during training, since they're ambiguous enough that including them would introduce noisy gradients.

Hard Negative Mining

After matching, the class imbalance problem becomes stark: a typical training image with 3–5 objects will produce 10–20 positive anchors out of 8,732 total. Training on all negatives simultaneously would cause the loss to be completely dominated by "background" predictions.

Hard negative mining addresses this by selecting only the hardest negatives, the background anchors with the highest classification loss, rather than all negatives. The number of selected negatives is capped at 3× the number of positives:

Loss Functions

SSD uses a multi-task loss that jointly optimizes for box localization and class prediction. The total loss is a weighted sum of the two components:

Localization loss is computed only over positive anchors (those matched to a ground-truth box), using Smooth L1 loss on the encoded box offset targets:

Where the smooth L1 function transitions between quadratic and linear to reduce sensitivity to large outliers:

Classification loss applies softmax cross-entropy over both positive and hard-negative anchors:

The box regression targets use a normalized encoding so the model predicts offsets relative to each prior's center and dimensions:

Here is the core smooth L1 implementation, which includes a selectable reduction strategy to support both per-image and per-batch normalization modes:

1def smooth_l1_loss(predicted_values, target, beta, reduction="sum"):
2  difference = predicted_values - target
3  absolute_difference = tf.math.abs(difference)
4 
5  small_mask = absolute_difference < beta
6  large_mask  = tf.logical_not(small_mask)
7 
8  errors = tf.where(small_mask, 0.5 * (difference ** 2) / beta, tf.zeros_like(difference))
9  errors = tf.where(large_mask, absolute_difference - (0.5 * beta), errors)
10 
11  # Sum over the four box coordinates
12  errors = tf.reduce_sum(errors, axis=-1)
13 
14  if reduction == "sum":
15      return tf.reduce_sum(errors)
16  elif reduction == "mean":
17      return tf.reduce_mean(errors)
18  else:
19      return errors  # per-anchor losses for hard negative selection

Mixed Precision Training (AMP)

Training in FP16 roughly doubles throughput on modern GPUs by fitting more data into VRAM and using the GPU's tensor cores. The challenge is that FP16's limited dynamic range (~65,504 max) can cause gradient underflow during backprop.

The AMPContext class handles this by wrapping the optimizer in TensorFlow's LossScaleOptimizer, which dynamically scales the loss upward before backprop and unscales gradients before the weight update. If scaling causes overflow, the update is skipped and the scale factor is reduced.

1class AMPContext:
2  def setup_policy(self):
3      if self._enabled:
4          policy = tf.keras.mixed_precision.Policy(self._policy)  # "mixed_float16"
5          tf.keras.mixed_precision.set_global_policy(policy)
6      else:
7          tf.keras.mixed_precision.set_global_policy("float32")
8 
9  def wrap_optimizer(self):
10      if not self._enabled:
11          return self._base_optimizer
12 
13      if self._loss_scale == "dynamic":
14          # Dynamic loss scaling: automatically adjusts the scale factor
15          self.optimizer = tf.keras.mixed_precision.LossScaleOptimizer(
16              self._base_optimizer
17          )
18      else:
19          # Fixed loss scale for reproducibility
20          self.optimizer = tf.keras.mixed_precision.LossScaleOptimizer(
21              self._base_optimizer, initial_scale=float(self._loss_scale)
22          )
23      return self.optimizer

Exponential Moving Average (EMA)

Training a model is noisy in nature. This causes a certain amount of variance in the weights which can lead to a degree of suboptimal generalization of the weights. The best way to imaging would be an object rocking back and forth on a spring. The weights are constantly moving since the model is learning and making adjustments. To stabalize this, a way to mitigate that would be to take a running average of the weights over time. This is the core of the Exponential Moving Average (EMA) technique. I will go more in depth in a future post about the EMA technique and the decisions that led me to implementing it in the training pipeline as an independent module. The formula for updating the EMA weights is as follows:

Where $d$ is the decay rate (typically 0.999). Early in training when weights are still far from optimal, a fixed decay rate would give too much weight to bad early estimates. The implementation uses an adjusted decay ramp that starts low and approaches the configured value as updates accumulate:

1def update(self, step: int):
2  if not self.should_update(step):
3      return
4 
5  decay = tf.constant(self._decay, tf.float32)
6  num_updates = tf.cast(self._num_updates, tf.float32)
7 
8  # Ramp up slowly at the start to avoid averaging in bad early weights
9  adjusted_decay = (1 + num_updates) / (10 + num_updates)
10  decay_rate = tf.minimum(decay, adjusted_decay)
11 
12  for ema_var, model_var in zip(self._ema_vars, self._model_vars):
13      d  = tf.cast(decay_rate, ema_var.dtype)
14      ema_var.assign(d * ema_var + (1.0 - d) * model_var)
15 
16  self._num_updates.assign_add(1)
17 
18@contextmanager
19def eval_context(self, model=None):
20  # Temporarily swap in EMA weights for evaluation, then restore training weights
21  use_ema = self.should_apply_during_eval()
22  if use_ema:
23      self.apply_to(model)
24  try:
25      yield
26  finally:
27      if use_ema:
28          self.restore(model)

The eval_context() context manager is the key interface: it swaps EMA weights into the model before evaluation and atomically restores the training weights afterward. If a SIGTERM arrives during evaluation, the finally block guarantees that weights are always restored correctly.

Learning Rate Scheduling

There is a key aspect to the model training that determines how well the model takes the steps in converging to the optimal weights. This would be the learning rate. The learning rate is a hyperparameter that is responsible for how much the weights are updated during the overall training cycle. The learning rate varies based on where the model is in the training cycle at a particular point in time.

A Learning Rate Scheduler follows a particular schedule that is defined to it at the beginning of the training cycle. The learning rate schedule goes through this predefined schedule and uses it to update the intensity behind the weight updates. It is similar to the human learning process. When starting with a new concept, the learning at that particular period is slow and needs to be warmed up. Once the concept starts to gain traction, the learning starts to crystallize and converge to a point where meaningful results can be observed.

Consequently, the learning rate schedule at the tail end of the training cycle does not need to make the model make large changes. The assumption is that by the end only some fine changes are required for the model to be ready and converge to the solution. This requires the learning rate schedule to have its phases such as "warmup" and "decay" phases.

In my pipeline I utilized a classic learning rate schedule that is a combination of Cosine Annealing and Linear Warmup. The easiest way to understand this is that in the beginning of the cycle the model is allowed to warm up by taking an increasingly large steps to lock on to a meaningful starting point. Once this is done, the model takes a gradually decreasing learning rate to converge onto its optimal weights. So by the end the changes are very small and the model is able to converge to a solution without any last minute overshooting or divergence. The formula for the learning rate schedule is as follows:

The warmup phase ramps learning rate from zero to base_lr over the first warmup_steps steps, which avoids large gradient updates at the start when batch statistics are unstable. After warmup, cosine decay gradually reduces the rate to min_lr over the remaining steps.

1class CosineWarmupSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
2  def __call__(self, step: tf.Tensor):
3      step = tf.cast(step, dtype=tf.float32)
4      warmup_steps = tf.cast(self.warmup_steps, dtype=tf.float32)
5      total_steps  = tf.cast(self.total_steps,  dtype=tf.float32)
6 
7      # Phase 1: linear warmup
8      warmup_lr = self.base_learning_rate * tf.minimum(
9          1.0, step / tf.maximum(1.0, warmup_steps)
10      )
11 
12      # Phase 2: cosine decay
13      progress   = tf.clip_by_value((step - warmup_steps) / (total_steps - warmup_steps), 0.0, 1.0)
14      cosine_lr  = self.minimum_learning_rate + 0.5 * (
15          self.base_learning_rate - self.minimum_learning_rate
16      ) * (1 + tf.math.cos(self.pi * progress))
17 
18      return tf.where(step < warmup_steps, warmup_lr, cosine_lr)

Results

All of the above is only worth doing if the predicted boxes actually converge. The clip below tracks a fixed validation frame across training checkpoints. Early epochs show wide, jittery, low-confidence boxes scattered around the object; by the later epochs the boxes have tightened around the true bounding box and confidence has stabilized.

The raw validation numbers are noisier than the video suggests. mAP at both IoU thresholds sees periodic drops toward zero throughout training, almost certainly from validation batches with few or no ground-truth boxes for a given class (see the "Avg GT Count per Image" and "Zero-GT Batch Ratio" diagnostics below). This only shows up in the VOC-computed metrics, however. COCO's mAP@0.5 never dips below ~0.76 across the entire run, which points at a quirk in how the VOC AP is averaged across classes rather than the model itself forgetting and re-learning detection dozens of times. mAP@0.50 climbs steadily and stays well ahead of the stricter mAP@0.75, with a brief dip for both around step 21k–29k before both recover and hold their plateau:

Conclusion