Implementation:NVIDIA DALI YOLOv4 Model
| Knowledge Sources | |
|---|---|
| Domains | Object_Detection, TensorFlow |
| Last Updated | 2026-02-08 16:00 GMT |
Overview
Implements the complete YOLOv4 object detection model in TensorFlow/Keras, including CSPDarknet53 backbone with SPP, PANet-style YOLO head, multi-scale loss computation, and Darknet weight loading.
Description
This module provides the `YOLOv4Model` class, a `tf.keras.Model` subclass implementing the YOLOv4 architecture for real-time object detection. The model consists of three main components: a CSPDarknet53 backbone with Spatial Pyramid Pooling (SPP), a Path Aggregation Network (PANet)-style YOLO detection head, and a multi-scale loss function.
The backbone (`CSPDarknet53WithSPP`) uses Cross Stage Partial (CSP) connections with Mish activation functions in residual blocks of increasing depth (1, 2, 8, 8, 4 repeats). It extracts three feature maps at different scales (route_1, route_2, and the final output) and applies SPP with max pooling at sizes 5, 9, and 13 to capture multi-scale context. The YOLO head (`YOLOHead`) uses upsample and downsample convolution blocks to create a feature pyramid, producing detection outputs at three scales (small, medium, large bounding boxes).
The `calc_loss` function computes the multi-scale detection loss combining: GIoU-based box regression loss (weighted 0.05), objectness loss with IoU-based masking to ignore predictions overlapping non-assigned ground truth (weighted 1.0), and classification loss using squared error against one-hot targets (weighted 0.5). Anchor assignment uses IoU matching between ground truth and predefined anchor sizes normalized to the 608x608 input resolution.
The model supports loading pretrained weights from Darknet-format binary files (`.weights`) by parsing the sequential layout of batch normalization parameters and convolutional kernels with proper dimension transposition. It provides built-in training step with loss tracking and learning rate logging, and a test step computing mAP via a Python function callback.
Usage
Use this model for YOLOv4-based object detection training within the DALI TensorFlow examples. Instantiate with the desired number of classes and input image size, optionally load pretrained Darknet weights, then train using standard Keras fit workflow.
Code Reference
Source Location
- Repository: NVIDIA_DALI
- File: docs/examples/use_cases/tensorflow/yolov4/src/model.py
- Lines: 1-340
Signature
def calc_loss(layer_id, gt, preds, debug=False) -> tf.Tensor:
"""Compute detection loss for one scale level."""
...
class YOLOv4Model(tf.keras.Model):
def __init__(self, classes_num=80, image_size=(608, 608)): ...
def fit(self, dataset, **kwargs): ...
def train_step(self, data) -> dict: ...
def test_step(self, data) -> dict: ...
def load_weights(self, weights_file): ...
# Architecture building methods
def darknetConv(self, filters, size, strides=1, batch_norm=True,
activate=True, activation="leaky"): ...
def darknetResidualBlock(self, filters, repeats=1, initial=False): ...
def CSPDarknet53WithSPP(self): ...
def yoloUpsampleConvBlock(self, filters): ...
def yoloDownsampleConvBlock(self, filters): ...
def yoloBboxConvBlock(self, filters): ...
def YOLOHead(self): ...
Import
from src.model import YOLOv4Model
model = YOLOv4Model(classes_num=80, image_size=(608, 608))
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| classes_num | int | No | Number of object detection classes (default: 80 for COCO) |
| image_size | tuple | No | Input image dimensions as (height, width); default: (608, 608) |
| input | tf.Tensor | Yes | Input image tensor of shape [batch, height, width, 3] |
| gt_boxes | tf.Tensor | Yes | Ground truth boxes and labels tensor [batch, max_objects, 5] with (x, y, w, h, class) |
| weights_file | str | No | Path to Darknet `.weights` file or Keras `.h5` file |
Outputs
| Name | Type | Description |
|---|---|---|
| small_bbox | tf.Tensor | Small-scale detection output [batch, grid_h, grid_w, 3*(classes+5)] |
| medium_bbox | tf.Tensor | Medium-scale detection output [batch, grid_h, grid_w, 3*(classes+5)] |
| large_bbox | tf.Tensor | Large-scale detection output [batch, grid_h, grid_w, 3*(classes+5)] |
| train metrics | dict | Dictionary with 'loss' and 'lr' during training |
| eval metrics | dict | Dictionary with 'mAP' during evaluation |
Usage Examples
Train YOLOv4 on COCO
from src.model import YOLOv4Model
# Create model
model = YOLOv4Model(classes_num=80, image_size=(608, 608))
# Load pretrained Darknet weights
model.load_weights("yolov4.weights")
# Compile and train
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-4))
model.fit(
train_dataset,
epochs=300,
steps_per_epoch=steps_per_epoch,
initial_epoch=0,
validation_data=val_dataset,
)