This repository presents a research project developed for learning purposes and voluntary participation in the Global Wheat Head Detection 2021 Challenge, which has already been concluded. Although the competition is no longer active, the original evaluation protocol was preserved to enable a fair comparison with previous approaches. Therefore, all experiments reported in this repository were evaluated using the Average Domain Accuracy (ADA) metric, the official metric adopted by the challenge.
The primary objective of this work was to investigate Knowledge Distillation (KD) as a model compression technique for object detection. This topic is particularly relevant for Precision Agriculture, where deep learning models are frequently deployed on embedded computing platforms with limited computational resources, memory, and power consumption. In such scenarios, compact models capable of maintaining competitive accuracy while significantly reducing computational cost are essential for real-time field applications.
To investigate this problem, baseline YOLOv5 models with different sizes (YOLOv5n, YOLOv5s, and YOLOv5m) were trained and evaluated. Subsequently, three teacher-student Knowledge Distillation configurations were investigated to determine how effectively knowledge from larger models could be transferred to lightweight detectors.
The entire experimental workflow was managed using MLflow, including experiment tracking, hyperparameter logging, metric monitoring, model versioning, artifact management, checkpoint storage, and reproducibility. This ensured that every experiment could be fully reproduced while maintaining complete traceability of datasets, training configurations, generated models, and evaluation results.
Two training datasets were used throughout the experiments:
| Dataset | Description |
|---|---|
| set0 | Original dataset proposed by the Global Wheat Head Detection 2021 Challenge. https://huggingface.co/datasets/claytonsds/gwhd2021_set0 |
| set1 | Augmented version of the original dataset created for this work. https://huggingface.co/datasets/claytonsds/gwhd2021_augmentedset1 |
The set0 dataset corresponds to the original training split proposed by the challenge. The set1 dataset was generated by applying an extensive data augmentation pipeline to improve the robustness and generalization capability of the detection models.
The augmentation pipeline was implemented using Albumentations 2.0.8 and included:
- Horizontal Flip (50%)
- Vertical Flip (20%)
- Random Affine Transformations
- Rotation (±15°)
- Scaling (0.8–1.2×)
- Shearing (±10°)
- Translation (10%)
- Random Brightness and Contrast
- Motion Blur or Gaussian Blur
- Color Jitter
- Random Shadow
- Coarse Dropout (Random Occlusions)
- JPEG Compression
- Image resizing to 640 × 640
All transformations preserved the object bounding boxes using the Pascal VOC format with a minimum visibility threshold of 30%.
All models were evaluated on the official test set of the Global Wheat Head Detection 2021 Challenge, following the original benchmark protocol and using the Average Domain Accuracy (ADA) metric.
The official test set contains 1,382 images distributed across 18 independent domains (see Figure 1), ensuring that the evaluated models are assessed on diverse environmental conditions, acquisition devices, wheat cultivars, and geographical locations.
The ADA metric computes the average detection accuracy across all domains rather than over all images, encouraging models that generalize well to unseen acquisition conditions instead of overfitting to the largest domains.
All experiments were evaluated using the Average Domain Accuracy (ADA) metric following the official evaluation protocol of the Global Wheat Head Detection 2021 Challenge.
The experiments demonstrate that training with the augmented set1 dataset substantially improves detection performance compared to the original dataset (see Figure 2). For example, the baseline YOLOv5s increased from 45.70% ADA using the original set0 dataset to 58.49% ADA when trained with the augmented set1, highlighting the effectiveness of the proposed augmentation strategy.
 |
 |
 |
 |
 |
 |
 |
 |
Knowledge Distillation further improved the lightweight student models. The best distilled model (YOLOv5m → YOLOv5s) achieved an ADA of 58.63% (see Figure 3), representing a 19.91% relative improvement over the corresponding student baseline while reaching performance very close to the larger YOLOv5m teacher (59.93% ADA).
| Model | Training Dataset | ADA Score | Download |
|---|---|---|---|
| YOLOv5m | set1 | 59.93% |  HF |
| YOLOv5s* | set0 | 45.70% | HF |
| YOLOv5s | set1 | 58.49% | HF |
| YOLOv5n | set1 | 44.75% | HF |
| YOLOv5s → YOLOv5n | set1 | 47.93% (+7.10%) | HF |
| YOLOv5m → YOLOv5s | set1 | 58.63% (-2.16%) | HF |
| YOLOv5m → YOLOv5n | set1 | 53.66% (+19.91%) | HF |
Although the largest model still achieved the highest overall accuracy, the distilled lightweight models provide a considerably better trade-off between detection accuracy and computational cost. Lightweight YOLOv5 architectures require significantly fewer parameters, lower memory consumption, and faster inference than larger models, making them especially suitable for embedded computer vision systems commonly employed in precision agriculture, such as UAVs, autonomous agricultural robots, and edge AI devices. These results demonstrate that Knowledge Distillation is an effective strategy for improving compact object detectors while maintaining performance close to much larger networks.
* YOLOv5s* trained on set0 corresponds to the original dataset proposed by the Global Wheat Head Detection 2021 Challenge and is reported to provide a direct comparison with the augmented training dataset (set1).
The proposed framework provides a complete and reproducible pipeline for training YOLOv5-based wheat head detectors. It integrates dataset preparation, Albumentations-based data augmentation, baseline model training, Knowledge Distillation, experiment tracking with MLflow, and final evaluation into a unified workflow.
The augmentation strategy generates two additional training samples and one additional validation sample for each original image, resulting in 18,275 training images and 2,952 validation images. Three baseline models (YOLOv5n, YOLOv5s, and YOLOv5m) were first trained independently, followed by three Knowledge Distillation experiments (YOLOv5s → YOLOv5n, YOLOv5m → YOLOv5n, and YOLOv5m → YOLOv5s). Throughout training, MLflow automatically records hyperparameters, metrics, checkpoints, visual artifacts, trained models, and experiment metadata, ensuring full reproducibility and traceability. Model performance is finally assessed using the Average Domain Accuracy (ADA) protocol of the Global Wheat Head Detection 2021 Challenge.
Figure 4 illustrates the complete scientific training pipeline adopted in this work.
The proposed Knowledge Distillation framework enables lightweight YOLOv5 models to learn from larger teacher networks through both prediction-level and feature-level supervision. During each training iteration, the input batch is simultaneously processed by the teacher and student networks. The teacher performs inference in evaluation mode (torch.no_grad()), while the student performs the standard forward and backward propagation. Intermediate feature maps are collected from both models using feature hooks for feature-level supervision.
A foreground mask is generated from the teacher objectness confidence, allowing only informative predictions to contribute to the distillation process. Four complementary losses are computed:
- Bounding Box Distillation using Smooth L1 Loss.
- Objectness Distillation using Binary Cross-Entropy (BCE).
- Classification Distillation using KL Divergence with temperature scaling.
- Feature Distillation using Mean Squared Error (MSE) between normalized attention maps.
The individual losses are combined into a weighted Knowledge Distillation objective, which is added to the original YOLOv5 detection loss. During optimization, only the student model is updated, while the teacher network remains frozen.
The loss formulation implemented in this repository is summarized below.
L_box = (1/N) Σ SmoothL1(student_box, teacher_box)
L_obj = (1/N) Σ BCEWithLogits(student_obj, sigmoid(teacher_obj))
L_cls = (T²/N) Σ KLDiv(
log_softmax(student_cls / T),
softmax(teacher_cls / T)
)
where T denotes the distillation temperature.
The attention map of each feature tensor is computed as
A(F) = mean(F², channel dimension)
and optimized using
L_feat = (1/L) Σ MSE(
A(student_features),
A(teacher_features)
)
L_KD =
λ_box · L_box
+ λ_obj · L_obj
+ λ_cls · L_cls
+ λ_feat · L_feat
L_total = L_YOLO + λ_KD · L_KD
where λ_KD corresponds to the kd_gain hyperparameter in the implementation.
Figure 5 illustrates the complete Knowledge Distillation pipeline implemented in this work, including the teacher and student forward passes, feature extraction, foreground masking, individual loss computation, and final loss aggregation.
