KD-LightNet: A Lightweight Network Based on Knowledge Distillation for Industrial Defect Detection

Abstract

At present, the method based on deep learning performs well in public object detection tasks. However, there are still two problems to be solved for industrial defect detection: 1) Industrial scenes requires real-time and lightweight; 2) Lightweight network accuracy is limited. In order to tackle these issues, based on knowledge distillation, this paper proposes an effective lightweight defect detection network (KD-LightNet) suitable for edge scene. First of all, a lightweight network (LightNet) is designed based on structure reparameterization, which can sufficiently improve the capability of network feature extraction and reduce the complexity of model inferring. Moreover, a well prepared self-distillation strategy is proposed, which utilize the pre-trained LightNet network as a teacher model to transfer knowledge in the same structure. Then, in order to fully utilize the logits predicted by teacher model, an improved KL divergence loss is proposed to enhance the accuracy of the student model. Finally, in the experiments, three industrial datasets (PKU-Market-PCB, NEU-DET and Magnetic tile defect dataset) were used to validate the proposed model performance. The KD-LightNet detection accuracy (mAP) is improved by an average of 6.87%, while the average detection speed reaches 72 FPS @3070Ti (Params 4.7M), which meets the requirements of industrial defect detection accuracy and real-time.