The paper introduces a novel methodology for reducing the size of YOLO-based underwater object detectors to deploy it on edge hardware. Feature extraction layer light-weighting technique is used to compress the model with minimal impact on performance. Two new object detection network topologies are created, suitable to be used as a student network in knowledge distillation tasks. Knowledge distillation algorithm with temperature decay strategy is developed to mitigate the performance loss caused by model compression without inflating the parameter count. Object detection models, based on the proposed methodology, are tested on Underwater Target Detection Algorithm Competition 2020 dataset, providing higher accuracy and offering faster runtime than the existing solutions.
Due to rapid development of deep learning [
To operate in real-time mode, the artificial intelligence model is usually deployed on a separate device, typically a single board computer, which is then installed on the AUV. This approach allows for greater modularity in system design and easier upgrades. However, this method has its own set of challenges and limitations. The edge devices chosen for this purpose are selected primarily for their economic and power efficiency, which is crucial for underwater operations. In such cases, performance is frequently traded off for extended battery runtime and lower overall system costs, potentially limiting the capabilities of installed software.
Another major problem is that detecting targets in underwater environments is significantly more challenging than generic object detection on land, due to the overall lower image quality of underwater datasets. The degradation in image quality is caused by presence of suspended particles in the water, which introduces significant color distortion and cast. One more issue, which further complicates the task, is motion blur, caused by continuous movement of underwater vehicle. Additionally, there are substantial differences in light emission and propagation between underwater and land scenes.
To overcome these obstacles, researchers have developed a variety of advanced techniques. Common approaches typically involve implementing the series of pre-processing steps to enhance image quality before inputting it into intellectual system, designing and deploying deeper neural networks capable of extracting meaningful features from degraded inputs, and integrating specialized network blocks that enhance feature representation abilities. Often, combinations of these methods are used, which results in complex, multi-stage processing pipelines. While these solutions have demonstrated impressive results in improving underwater object detection, they come with a significant drawback: increased overall network complexity, which translates directly into higher hardware requirements, larger model sizes, and increased power consumption.
The complexity of these advanced models makes it infeasible to deploy them on the edge devices typically used in AUVs. Thus, there is a need for research focused on size reduction strategies specifically tailored for underwater object detectors. The ultimate goal of such research should be to develop a comprehensive methodology that enables significant reductions in the size of object detection neural networks while simultaneously preserving their accuracy and operational speed.
The successful development of such methodologies could potentially lead to the design of smaller, more efficient AUVs capable of accessing environments that are currently out of reach. Considering this, this research paper aims to explore novel approaches to network size reduction for underwater object detection models, addressing this crucial challenge and contributing to the ongoing advancement of AUV technologies and their applications in marine science, environmental monitoring, and underwater operations.
Generic object detection neural network topologies development started in 2014, when R-CNN model
[
Later in 2015, one-stage object detection algorithms such as Single-Shot Detector (SSD) and You
Only Look Once (YOLO) [
To improve both speed and detection accuracy, YOLO series of object detectors have undergone
a lot of updates. YOLOv2 [
Generally, modern versions of YOLO network share the same topology, where convolutional neural network (backbone) is responsible for feature extraction at multiple scales, feature pyramid network (neck) is used to aggregate multi-scale features, mixing contextual and detailed information, and detection head is responsible for regression and classification tasks.
To overcome constraints presented by lower quality of underwater images, unique features of underwater targets such as small size and dense location, and computational constraints, specific underwater detection frameworks are developed. Typically, these frameworks can be divided into three categories by the way they achieve efficiency improvements for detecting targets in underwater environments.
The first category are object detectors with higher feature representation abilities, reached by
bigger network capacity or introduction of specific blocks for better feature extraction. Attention
mechanisms are often used to enhance feature extraction capabilities of the model by ensuring that
backbone layers of the network focus on more relevant features [
The second category of underwater object detection frameworks includes models, which achieve
accuracy gains by focusing on mitigating a specific issue with underwater images, such as small
target size, target overlap and motion blur [
The third category includes underwater object detection methodologies, which focus on
decreasing the parameter count, size and latency of the model [
This fact has sparked our interest in comparing the performance of light-weight underwater object detectors and finding optimal way to mitigate the performance loss, which is inevitably a sideeffect of model light-weighting.
Knowledge distillation (KD) is a highly efficient technique of boosting efficiency of a light-weight
student model by forcing it to mimic the outputs of a larger pre-trained teacher model. Earliest
overviews on these methodologies were formulated by Bucilua [
While being an efficient technique, KD process ensures best results when teacher and student models share the similar architecture (e.g. CNN-to-CNN, ViT-to-ViT). As feature heterogeneity increases in later layers of neural network, it becomes harder for a student model to improve performance. Researches by Touvron et al. [35], Hao et al. [36] are offering more efficient KD algorithms for heterogenous architectures, however, the problem is still not fully resolved and KD performance is better when teacher and student models are similar.
The problem of creating an optimal size reduction methodology for YOLO-based underwater object detectors involves building two neural networks: samples ( 1, 1), ( 2, 2), … , ( , ), where denotes the -th RGB image matrix of dimensions × × 3, and represents the vector of ground truth bounding box coordinates, which represent the position of the object within the image, and class labels for each object within the image. The teacher network parametrized by , parametrized by the weights , and the student network , are trained to handle the following transformation: : ↦ {( ̂ where ̂ is the predicted bounding box for object in image , and ̂ is the predicted class probability for the object. The goal is to minimize the discrepancy between the predictions ̂ and the ground truth , for bounding boxes and class labels, respectively.
The loss function of
, denoted as ℒYOLO, is composed of regression and classification losses. CIoU [37] is used as a main bounding box regression loss, defined as: ℒCIoU( , ̂)= 1 − IoU( , ̂)+ 2( , ̂ ) 2 + ⋅ , where , ̂ are the ground truth and predicted bounding boxes, respectively, IoU( , ̂) is the Intersection over Union, which measures the overlap between the predicted and ground truth bounding boxes:
where | ∩ ̂| is the area of overlap, | ∪ ̂| is the total area covered by both boxes, ( , ̂ )is the Euclidean distance between the centroids , ̂ , is the weight factor, balancing the importance of aspect ratio consistency:
IoU( , ̂)= | ∩ ̂| | ∪ ̂|
, =
(1 − IoU( , ̂))+ measures the consistency of the aspect ratio between the predicted and ground truth boxes: 4 2 ℎ = (arctan ( )− arctan ( )) , ̂ ̂ ℎ 2 where and ℎ are the width and height of the ground truth bounding box, and ̂ , ℎ̂ are the width and height of the predicted bounding box.
Classification is governed by varifocal loss function (VFL), defined as: ℒVFL(p, q)= { − ( − and is the penalty coefficient [38]. loss: where, is the predicted classification score, is the target score, is the balancing coefficient, Total loss function can be defined as the sum of these two loss functions and distribution focal ℒYOLO = CIoUℒCIoU + DFLℒDFL + VFLℒVFL, with CIoU, DFL and VFL being hyperparameters, balancing the importance of each component. To reduce the size of
while maintaining performance, knowledge distillation is used for transferring knowledge from the larger teacher model . The objective of knowledge distillation is to minimize a combination of the standard YOLO loss ℒYOLO and distillation loss ℒKD, for the student model is defined as follows: where controls the trade-off between the YOLO loss and the distillation loss in the total loss function.
Both with the following weight update rule: respect to the model parameters.
ℒstudent = ℒYOLO + (1 − )ℒKD,
+1 = − ∇ ℒ, where is the learning rate controlling step size, and ∇ ℒ is the gradient of the loss function with
To enable knowledge transfer from larger teacher model into light-weight student model, regression and classification components of total loss function have been modified to include distillation loss with additional weighting coefficient added to avoid learning collapse due to student model fully mimicking teacher model outputs. Additionally, temperature coefficient with decay strategy has been used to control the softening of logits for the classification component, allowing to regulate the amount of knowledge being distilled from a larger model by a student.
Bounding box regression loss is handled by CIoU function with additional L2 loss component, improving the consistency between teacher and student model bounding box predictions: ℒCIoU = (1 − distill)⋅ (1 − CIoU( , ̂ ))+ distill ⋅ ∥∥ ̂teacher − ̂student∥∥2, where distill is the distillation coefficient balancing the standard CIoU loss and the distillation term, ̂teacher and ̂student are the bounding boxes predicted by the teacher and student networks in that order.
Classification loss is handled with varifocal loss function (VFL), with Kullback-Leibler divergence loss (KL-loss) added as a distillation component, scaled by temperature. To encourage exploration in earlier stages of student model training, we propose temperature decay strategy, starting from higher values and linearly shifting the temperature coefficient toward 1 to focus on more confident predictions in later stages of training: ℒVFL = (1 − distill)⋅ VFL( ,
̂student)+ distill ⋅ ( )2 ⋅ KL( (
) ∥ ( ̂teacher ( ) ̂student ( ) )), (11) with ( )representing the decaying temperature, ̂teacher and ̂student are the class probability logits for teacher and student models, respectively.
Algorithm 1 Training YOLO student and teacher models
, learning rate , batch size . initial temperature = 0.
For each epoch = 1, … ,
: shuffle training set ; for each mini-batch ℬ ⊂ : Input: Dataset
= {( , )} =1, teacher model , student model , initial temperature 0, decay strategy ( ), hyperparameters CIoU, DFL, VFL, distill, maximum epochs Initialize student model parameters , load teacher model with parameters , set do forward path, computing ̂ compute loss ℒtotal; do backward pass, compute gradients ∇ update student ← student − ⋅ ∇
ℒ update temperature ( ).
To ensure optimal performance of a resulting distilled model, the student model should meet the size and computational efficiency requirement. An approach used in this paper involved light-weighting the feature extraction (backbone) layers of YOLO object detector, to reduce the number of expensive convolutional operations, which contribute a lot to a total parameter count. Feature aggregation (neck) and final output layers (head) from original YOLOv8 architecture were reused, as adding additional blocks to these parts of the network would increase the parameter count, and extra lightweighting would introduce more differences between student and teacher model, which could harm the distillation performance. In this case, backbone network consists of convolutional blocks, which are composed of 2D convolutional layer and batch normalization, followed by SiLU activation function. Each convolutional block is followed by bottleneck blocks C2f or C3, which perform convolutional operation on the input, then splits the channels, processes the resulting feature map through multiple bottlenecks (number in the name of the block represents the bottleneck layer count), ending with concatenation.
Overall architecture of object detection network is shown in Figure 1.
To find the model, which provides the best balance between efficiency and performance, we have built two variants of backbone network, based on GhostConvolution layers derived from GhostNet and FasterNet blocks. The core idea behind this design is to reduce the number of computationally expensive 3 x 3 convolutions in favor of their light-weight counterparts.
Default convolution operation outputs feature map by processing input feature map using a convolution kernel as = ∗ , where ∗ denotes the convolution operation. GhostConvolution is aimed at generating the similar number of feature maps by using less computations by executing convolution with fewer filters, obtaining intrinsic feature maps int as int = ∗ int where int is a smaller convolutional filter. Then, the series of computationally inexpensive operations is used to generate the ghost feature maps ghost as in ghost = ( int)where denotes cheap operation.
FasterNet applies different approach to reduce the computational complexity and decrease latency of convolutional operations, based on PConv procedure, which applies convolutional operation only on a part of input channels for spatial feature extraction and leaves remaining channels as is. Then, PConv is followed by series of pointwise convolution to reuse the information from all channels in an efficient way.
To evaluate the performance of the proposed algorithm, we used the UTDAC2020 (Underwater Target Detection Algorithm Competition 2020) dataset, a challenging underwater detection benchmark, consisting of 5168 training and 1293 validation images in various resolutions (3840 x 2160, 1920 x 1080, 720 x 405, and 586 x 480). The chosen dataset contains four classes: echinus, holothurian, scallop, and starfish. UTDAC2020 presents several challenges, including significant class imbalance, with the echinus class appearing four times more frequently than the other classes. The dataset also features targets at different scales, often densely packed, and the cases of low contrast, challenging lighting conditions and motion blur caused by camera movement.
The machine used for experiment is equipped with Intel Core i5-13600K processor, NVIDIA A4000 GPU with 16GB VRAM. Software-wise, the test setup is running Ubuntu 20.04.6 LTS with Python 3.10.13, CUDA 12.1, and PyTorch 2.2.1.
Each model has been trained during 250 epochs with a batch size of 32 and an image size of 640 x 640. Stochastic Gradient Descent (SGD) served as the optimization algorithm, with a momentum of 0.937, an initial learning rate of 0.01, and a weight decay coefficient of 0.005. For distilled models, initial temperature coefficient is set to 5, linearly decaying toward 1 until the model converges. The hyperparameters were found empirically using the Ray Tune library. Albumentations package was employed for augmenting the dataset, which includes combinations of random crop, random rotate and mosaic augmentations.
Total of five metrics have been used to test the model, with mAp and mAp50 representing the object detection accuracy of neural network. Size, parameter count and FLOPs are also measured as performance metrics to evaluate the computational efficiency of proposed approach. Detailed explanation of metrics is provided in Table 1.
Mean average precision (mAp) at intersection over union (IoU) of 0.50, defined as average precision for each class over number of classes mAp at IoU of 0.50:0.05:0.95 Total number of model parameters
Performance metrics denoting number of floating-point operations per second
Model size in megabytes
A comparison of existing frameworks and models, built using the proposed methodology, is provided in Table 2. For distilled models, YOLOv8l with DarkNet-53 backbone is used as a teacher model. Student models use YOLOv8s architecture with custom backbones, based on GhostNet and FasterNet, with both convolutional blocks and bottlenecks modified. Models using knowledge
Experiment results prove superior performance of models using knowledge distillation algorithm. In comparison to YOLOv8s, the computational complexity in FLOPs is reduced by 45%, while 63 maintaining object detection performance. The choice of light-weight provide any significant difference, with FasterNet and GhostNet backbones providing similar results in terms of both size and accuracy.
Visualization of ground truth bounding boxes and class labels, compared with detections, performed with the proposed YOLOv8-dist model, is shown in Figure 2. The samples with complex background information and targets of various sizes were selected to demonstrate the performance of our model.
The paper proposed a novel methodology to reduce the size of YOLO-based underwater object detectors. Knowledge distillation algorithm with temperature decay strategy has been developed for object detection neural network, allowing to effectively train light-weight student model by transferring knowledge from teacher model of larger capacity. Additionally, we developed two lightweight YOLO architectures, derived from GhostNet and FasterNet approaches to convolution operation, which are suitable to be used as a student model in knowledge distillation tasks.
The proposed light-weight models are 45% more efficient in terms of computational complexity, compared to existing YOLOv8s model. After using our knowledge distillation algorithm, the performance of the student model is superior to original YOLOv8s in terms of accuracy (84.72% and 84.58%, respectively), while the model size is comparable to YOLOv8n, the smallest model among YOLO-based detectors.
The applications of our methodology include training efficient object detection neural network for integrated autonomous underwater vehicle hardware. To achieve even better performance of knowledge distillation with YOLO-based detectors, we suggest that further research could be conducted on more sophisticated knowledge distillation techniques during training, and using knowledge distillation algorithms with heterogenous backbones, such as visual transformers, which could potentially enrich intermediate feature maps with more context and semantic information, leading to a higher accuracy in applied tasks.
[34] Y. Zhang, T. Xiang, T. M. Hospedales, H. Lu, Deep Mutual Learning, in: 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), IEEE, 2018. doi:10.1109/cvpr.2018.00454. [35] H. Touvron, M. Cord, M. Douze, F. Massa, A. Sablayrolles, H. Jegou, Training data-efficient image transformers & distillation through attention, in: Proceedings of the 38th International Conference on Machine Learning, Journal of Machine Learning Research, 2021, pp. 10347 10357. URL: https://proceedings.mlr.press/v139/touvron21a.html. [36] Z. Hao, J. Guo, K. Han, Y. Tang, H. Hu, Y. Wang, C. Xu, One-for-All: Bridge the Gap Between Heterogeneous Architectures in Knowledge Distillation, 2023. URL: https://arxiv.org/abs/2310.19444. [37] Z. Zheng, P. Wang, D. Ren, W. Liu, R. Ye, Q. Hu, W. Zuo, Enhancing Geometric Factors in Model Learning and Inference for Object Detection and Instance Segmentation, 2020. URL: https://arxiv.org/abs/2005.03572. [38] H. Zhang, Y. Wang, F. Dayoub, N. Sunderhauf, VarifocalNet: An IoU-aware Dense Object Detector, in: 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), IEEE, 2021. doi:10.1109/cvpr46437.2021.00841.