A Lightweight Transfer Learning Model with Pruned and Distilled YOLOv5s to Identify Arc Magnet Surface Defects

Abstract

Surface defects in arc magnets constitute the main culprit for performance degradation and safety hazards in permanent magnet motors. Machine-vision methods offer the possibility to identify surface defects automatically. However, the current methods still do not adequately solve the problems of low identification accuracy, excessive dependency on training data, and sizeable computational complexity. This paper proposes a lightweight YOLOv5s-based transfer learning model with network pruning and knowledge distillation to address these issues. Our model was derived from a pre-trained YOLOv5s for general object detection. A transfer learning mechanism was designed to obtain the optimal surface defect identification accuracy of the model from fewer training samples. Network pruning and knowledge distillation were combined to compress the transferred model. The transferred model serves as the teacher model of knowledge distillation, while its pruned model acts as the student model. To weaken the loss of the accuracy after model compression, a new λ factor was introduced into the confidence loss function of the student model to increase the sensitivity of identifying the defects. The experimental results show that our model’s performance is higher than other regular lightweight models. The identification accuracy for different defective arc magnets could reach 100%, the model size could achieve 1.921 MB, and the average inference time was 9.46 ms. Our model also has high accuracy in other defect identification applications besides arc magnets.