Grabbing detection training method based on mixed semi-supervised and knowledge distillation

Abstract

The invention discloses a hybrid semi-supervised and knowledge distillation-based capture detection training method. The method comprises the steps of constructing a labeled data set and an unlabeled data set; the method comprises the following steps: performing supervised training on a pre-constructed first network model for capturing detection by using a labeled data set, and detecting an enhanced unlabeled data set by using a captured detection model obtained by training to obtain a pseudo label of the unlabeled data set; obtaining a mixed data set from the labeled data set and the data set with the pseudo label, and performing semi-supervised training on the initially pre-trained capture detection model; and taking the grabbing detection model obtained through semi-supervised training as a teacher model, and performing guided training on the lightweight student network model by adopting a knowledge distillation method and based on the mixed data set to obtain a lightweight grabbing detection model. According to the method, the high-precision lightweight grabbing detection model can be trained at low labeling cost.

Description

Technical Field

[0001] The present invention belongs to the technical field of intelligent grasping detection, and in particular relates to a grasping detection training method based on hybrid semi-supervision and knowledge distillation. Background Art

[0002] With the development of computer technology, grasp detection methods based on deep learning have been widely used. Methods based on deep learning are mainly divided into discriminative methods and generative methods. Discriminative grasp detection methods rely on the region-based convolutional neural network (RCNN) structure. The model combines the grasp region proposal network to generate candidate regions for grasp detection, and further filters out the optimal grasp posture on this basis. The generative strategy directly predicts the grasp posture without detecting the candidate grasp posture. It relies on the full convolutional neural network. The generative method is usually a single stage, so it has a significant advantage over the discriminative strategy in terms of algorithm efficiency. Some later studies adopted a pixel-level generative method similar to image segmentation based on the fully connected neural network (FCNN). The output of this method is usually represented by a set of heat maps, which represent the quality, angle and width of the grasp rectangle. Compared with the previous generative strategy, the proposed pixel-level method makes grasp detection more fine-grained and accurate. These pixel-level generative methods achieve the same level of accuracy as the discriminative strategy and are more efficient. Existing grasp detection methods find it difficult to strike a good balance between accuracy and lightweight, and often perform poorly in application scenarios with limited computing resources. Therefore, we propose a generative lightweight grasp detection network.

[0003] Grasp detection methods based on supervised learning require a large number of labeled data sets, and the labeling of grasp data sets is a time-consuming and labor-intensive process. The generalization ability of the model often decreases when facing different application scenarios. Therefore, researchers have proposed grasp detection methods based on semi-supervised learning. Semi-supervised learning methods in the field of grasp detection mainly include variational auto-encoder (VAE) encoding and confidence-driven methods based on uncertainty estimation and geometric consistency. Vector Quantized Variational AutoEncoder (VQ-VAE) encodes high-dimensional input data into low-dimensional discrete representations to achieve data compression. The encoder compresses the input data into a smaller feature vector, and the quantization step converts the feature vector into a discrete index representation. The confidence-driven method based on uncertainty estimation largely prevents the student model from learning wrong and harmful information from consistency loss, accelerates the learning progress, and further improves the accuracy of the model. None of these methods involve the use of deep information, and there is still room for improvement in the accuracy of grasp detection. Summary of the invention

[0004] The present invention provides a grasping detection training method based on hybrid semi-supervision and knowledge distillation, which trains a high-precision lightweight grasping detection model with low annotation cost.

[0005] In order to achieve the above technical objectives, the present invention adopts the following technical solutions: A grasp detection training method based on hybrid semi-supervision and knowledge distillation, including: Obtain a training data set, each training data includes an RGB image and a depth feature map of the same captured area; Mark the crawling configuration of some training data in the training data set to build a labeled data set , the remaining training data is used to construct an unlabeled dataset ; For labeled datasets and unlabeled datasets Perform data enhancement to obtain enhanced labeled data sets and enhanced unlabeled datasets; Using labeled datasets Performing supervised training on a first network model pre-built for grasping detection to obtain an initial pre-trained grasping detection model; Use the initial pre-trained grasp detection model to detect the enhanced unlabeled dataset to obtain an enhanced dataset with pseudo labels. ; Enhanced labeled dataset With enhanced pseudo-labeled datasets Mix to get a mixed data set ; Using a mixed dataset , semi-supervised training of the initial pre-trained grasp detection model; The grasp detection model obtained by semi-supervised training is used as the teacher model, and the knowledge distillation method is used based on the mixed dataset. , guiding the training of a second network model used for grasping detection; wherein, the second network model is a lightweight network obtained by simplifying the first network model.

[0006] Furthermore, the first network model for grasping detection includes an encoder stage, a residual layer stage and a decoder stage; the encoder stage includes a plurality of convolutional layer modules and a plurality of fusion modules respectively following each convolutional layer module, the residual layer stage includes a plurality of residual modules, and the decoder stage includes a plurality of convolutional transposition layers; The second network model for grasping detection is based on the first network model, with the number of fusion modules, residual modules and convolutional transpose layers reduced.

[0007] Furthermore, the fusion module is an RGB-D fusion module that combines channel attention, spatial attention and channel exchange; in the fusion module: First, the channel attention module processes the RGB image and the depth feature map from the same area and with the same shape to obtain the weight of the channel dimension; Then, the channel exchange module uses the weight information obtained by channel attention to perform channel exchange operations on the RGB image and depth feature map of the input fusion module; The channel swap operation is as follows: if the weight of a feature map in the cth channel is lower than a certain threshold, and the weight of another feature map in the cth channel exceeds another threshold, the cth channel of the former feature map is replaced by the cth channel of the latter feature map; Finally, the spatial attention module processes the feature map after channel swapping to obtain a new color-depth feature map.

[0008] Further, the residual module is described as: ; ; in, is the input of the residual module, represents convolution, represents regularization, represents global average pooling, and are the feature maps of different intermediate layers, and Both are activation functions; represents element-wise multiplication, is the output of the residual module.

[0009] Furthermore, the data enhancement performed on the data set includes performing random rotation, mirroring and / or symmetry operations on the images in each training data.

[0010] Furthermore, using a mixed dataset , when semi-supervised training is performed on the initial pre-trained grasp detection model, the loss function is expressed as: ; in, Initial pre-trained grasp detection model based on dataset The L1 smooth loss for training, Initial pre-trained grasp detection model based on dataset The L1 smooth loss for training, is the coefficient, is the total loss for semi-supervised training.

[0011] Furthermore, the loss function for training the second network model is expressed as: ; ; in, is the loss for training the second network model, For the student model based on the mixed dataset The L1 smooth loss between the predicted output and the label; Predict output for the teacher model The predicted output with the student model The L2 loss between them is used to minimize the feature difference between the teacher model and the student model; The input images of the teacher model and the student model are mixed data sets. image training data.

[0012] The lightweight grasping detection model trained by the present invention can fully integrate the modal information of RGB-D, effectively improve the grasping detection accuracy, and significantly improve the performance of the grasping detection model in the case of limited annotated data sets. Compared with the existing grasping inspection methods and training methods, it has the following advantages: (1) The modal fusion module designed in the present invention effectively improves the degree of information fusion and provides improved spatial and channel sensitivity to enhance model performance; the proposed residual module effectively alleviates the problem of gradient dissipation.

[0013] (2) On edge computing devices with limited computing resources and power consumption, the lightweight student model based on knowledge distillation proposed in the present invention performs model tailoring in the encoding and decoding stages based on the teacher model, which greatly reduces the number of model parameters while still maintaining a relatively fast grasping and detection speed, thus meeting real-time requirements.

[0014] (3) The training method based on hybrid semi-supervision enables the model to make full use of the data. The diversity of training samples is greatly improved through random enhancement and pseudo-labeling methods. It can still effectively improve the model accuracy when there are insufficient labeled samples, reduce the cost of manual labeling, and provide a potential solution for improving the generalization ability of the model.

[0015] (4) The grasping detection method based on modal fusion proposed in this paper has obvious advantages in grasping accuracy compared with other common grasping detection methods. On the Cornell dataset, the accuracy reached 98.5%; on the multi-target grasping dataset, the accuracy reached 95.4%; the model achieved a good balance between the number of parameters and performance. The student grasping detection model trained based on the hybrid semi-supervised and knowledge distillation method achieved higher grasping detection accuracy under the same training samples and training environment. BRIEF DESCRIPTION OF THE DRAWINGS

[0016] Figure 1 This is the first grasping detection network proposed by the present invention.

[0017] Figure 2 It is the modal fusion module of the present invention.

[0018] Figure 3 It is the exchange module in the modal fusion module.

[0019] Figure 4 An efficient residual module designed for the present invention.

[0020] Figure 5 Schematic diagram of the second stage semi-supervised training.

[0021] Figure 6 Schematic diagram of the third stage of knowledge distillation training.

[0022] Figure 7 This is a flow chart of the overall training strategy. DETAILED DESCRIPTION

[0023] The following is a detailed description of an embodiment of the present invention. This embodiment is based on the technical solution of the present invention, and provides a detailed implementation method and a specific operation process to further explain the technical solution of the present invention.

[0024] The embodiment of the present invention provides a method for training lightweight networks for grasping and detection based on hybrid semi-supervision and knowledge distillation for garbage collection scenarios with conveyor belts as the background. At the input end, the color image and the depth image are efficiently fused in the designed modal fusion module, and then the feature map passes through the residual module to alleviate gradient dissipation, and finally the grasping configuration is output; the training is divided into three stages, first supervised training, then data enhancement is performed on the basis of supervised training, semi-supervised training is performed, and finally the knowledge distillation method is used on the basis of semi-supervised training, combining the predictions of the teacher model and the student model to enhance the generalization ability of the student model, while improving the robustness and performance of the model, and finally training a high-precision lightweight grasping and detection model with low annotation cost.

[0025] The grasping detection training method based on hybrid semi-supervision and knowledge distillation in this embodiment includes the following steps: Step 1: Obtain a training data set, each training data set includes an RGB image and a depth feature map of the same captured area.

[0026] Step 2.1: Mark the crawling configuration of some training data in the training data set to build a labeled data set , the remaining training data is used to construct an unlabeled dataset .

[0027] , ; in, For labeled datasets There are n training data in the is the corresponding label; Unlabeled dataset There are m training data in .

[0028] Step 2.2: For the labeled dataset and unlabeled datasets Perform data enhancement to obtain enhanced labeled data sets and enhanced unlabeled datasets.

[0029] For the dataset and The data enhancement operation performed can be performed by randomly rotating (90 degrees, 180 degrees, 270 degrees), mirroring and / or symmetry of the images in each training data. , it is necessary to perform the same random data augmentation on the images and labels to obtain an enhanced labeled dataset .

[0030] Step 3: Use the labeled dataset The first network model pre-built for grasping detection is subjected to supervised training to obtain an initial pre-trained grasping detection model.

[0031] The first network model for grasping detection includes an encoder stage, a residual layer stage, and a decoder stage. The first network model of this embodiment is as follows Figure 1 As shown. In the input stage, the RGB image and depth image of size 224×224 are input into the network separately instead of being concatenated together. The encoder stage includes 3 convolutional layer modules and 3 fusion modules following each convolutional layer module. In order to simplify the network structure, the result of the last fusion module is concatenated along the channel dimension, and the output size is 256×56×56. The residual layer stage includes 5 residual modules to alleviate the gradient dissipation problem. The tensor size in the residual module is kept at 256×56×56. The decoder stage includes 3 convolutional transpose layers to restore the tensor size to 32×224×224. For the feature map output by the last fusion module, it is output to the residual module by splicing.

[0032] The fusion module, such as Figure 2 Shown is an RGB-D fusion module that combines channel attention, spatial attention, and channel swapping to fully utilize information from different modalities. The encoder stage includes 3 repeated convolutional layer modules and 3 fusion modules, which are applied after the convolutional layer and regularization function. In the fusion module: First, the channel attention module processes the RGB image and the depth feature map from the same area and with the same shape to obtain the weight of the channel dimension;

[0033] In this embodiment, the channel attention consists of a convolution layer with a 3×3 convolution kernel and a regularization function, a global average pooling layer, and then a normalization function. It can be described as: ; in, , are the output and input features of channel attention respectively, They are normalization function, global average pooling, regularization function, and convolution.

[0034] Then, the channel exchange module uses the weight information obtained by channel attention to perform channel exchange operations on the RGB image and depth feature map of the input fusion module; Among them, Figure 3 As shown in , the channel swap operation is as follows: if the weight of a feature map in the cth channel is lower than a certain threshold, and the weight of another feature map in the cth channel exceeds another threshold, then the cth channel of the former feature map is replaced by the cth channel of the latter feature map. The condition can be described as: ; in, , is the weight extracted by channel attention, To specify the threshold, the subscripts represent different modes. In this embodiment, the threshold is set to the weight value corresponding to the 35% percentile position after the weights are sorted in ascending order.

[0035] Finally, the spatial attention module processes the channel-swapped feature map to enhance the spatial sensitivity of the feature and obtain a new color-depth feature map.

[0036] Spatial attention first captures the mean and maximum of the input feature map along the spatial dimension, and then concatenates these values ​​using a 3×3 convolution kernel before the convolution layer, and uses regularization and normalization functions after the convolution layer. The output of spatial attention can be described as: ; in, , are the output and input features of spatial attention, respectively. They are normalization function, global average pooling, regularization function, convolution, average and maximum operations.

[0037] In addition, this embodiment also designs an efficient residual module, such as Figure 4 As shown in Figure 1. First, two convolutional layers with 3×3 convolution kernels are used to extract basic features; regularization and activation functions are used after each convolutional layer. Then, a global average pooling layer and a convolutional layer are used before the normalization function to improve the channel sensitivity. The residual is obtained by multiplying the basic features with the channel weights. Finally, the input feature map and the residual are added element by element using the activation function. In this residual module, the output residual is described as: ; ; in, is the input of the residual module, represents convolution, represents regularization, represents global average pooling, and are the feature maps of different intermediate layers, and Both are activation functions; represents element-wise multiplication, is the output of the residual module.

[0038] Step 4: Use the grasping detection model obtained by the initial pre-training in step 3 to detect the enhanced unlabeled dataset to obtain an enhanced dataset with pseudo labels. ; Enhanced labeled dataset With enhanced pseudo-labeled datasets Mix to get a mixed data set , expressed as: ; Mixed Datasets It is several times the size of the original dataset, greatly increasing the diversity of the dataset and improving its distribution.

[0039] Step 5: Use a mixed dataset , semi-supervised training is performed on the initial pre-trained grasp detection model.

[0040] Using a mixed dataset , when semi-supervised training is performed on the initial pre-trained grasp detection model, the loss function is expressed as: ; in, is the total loss of semi-supervised training, Initial pre-trained grasp detection model based on dataset The L1 smooth loss for training, Initial pre-trained grasp detection model based on dataset The L1 smooth loss for training, is the coefficient, and in this embodiment, the value is 0.8. Figure 5 shown.

[0041] Step 6: Use the grasp detection model obtained by semi-supervised training in step 5 as the teacher model, adopt the knowledge distillation method and based on the mixed dataset , guiding the training of the second network model for grasp detection.

[0042] The second network model is a lightweight network obtained by simplifying the first network model.

[0043] The second network model used for grasping detection is based on the first network model, and the number of fusion modules, residual modules and convolutional transposition layers are deleted. The second network model of this embodiment is a lightweight model that removes one convolution operation in the encoder stage, two residual layers in the residual layer stage, and one convolutional transposition layer in the decoder stage as a student model. Since the lightweight student model has a small number of parameters and a simple model structure, the problem of decreased accuracy occurs in traditional supervised learning training methods. Therefore, the embodiment of the present invention adopts an offline knowledge distillation method, that is, a teacher model is used to predict the grasping detection task (generated offline), and then the prediction results of the teacher model are used as guidance for the student model, thereby improving the performance of the student model. Figure 6As shown, the loss function for training the second network model (used as the student model in the present invention) is expressed as: ; ; in, is the loss for training the second network model, For the student model based on the mixed dataset The L1 smooth loss between the predicted output and the label; Predict output for the teacher model The predicted output with the student model The L2 loss between the teacher model and the student model describes the square of the L2 norm (Euclidean distance) between the feature outputs of the teacher model and the student model, which is used to minimize the feature difference between the teacher model and the student model; The input images of the teacher model and the student model are mixed data sets. image training data.

[0044] The following provides an embodiment of the method of the present invention, which is experimentally or engineeringly applied using the Python programming language, comprising the following steps: Step 1: Verify the grasping detection network proposed by this method on the Cornell dataset and the multi-target solid waste dataset.

[0045] The training hyperparameters on the Cornell dataset are consistent with other related works. The batch size is set to 8, the learning rate is 0.001, and the optimizer is Adam. In the experiment, the grasp detection model proposed by the method of the present invention is compared with the common methods in the grasp detection field. The ratio of the training set to the validation set is 9:1, and the detection accuracy is shown in Table 1.

[0046] ; It can be found that the method of the present invention can achieve high grasping detection accuracy while having a smaller number of parameters.

[0047] On the multi-target solid waste dataset, the batch size is set to 16, the learning rate is 0.001, and the optimizer is Adam. In the experiment, compared with the common methods in the grasping detection field, the ratio of the training set to the validation set is 8:2, and the detection accuracy is shown in Table 2.

[0048] ; Among them, T1 refers to the inference speed on NVIDIA 4090, and T2 refers to the inference speed on the edge computing device Jetson Xaiver NX.

[0049] In this example, training requires Ubuntu 20.04 LTS and above, and the system environment requires Python 3.11.5, Pytroch 2.2.1, Torchvision 0.18.0, CUDA 12.1 and above. The hardware platform needs to meet the requirements of NVIDIA GeForce RTX 4090 (24G) graphics card, 32G or more memory and a hard disk with a capacity of no less than 256G. Use random seeds during training so that the same random splits are obtained when running training at different times, eliminating differences introduced by randomness.

[0050] Step 2: Verify the grasping detection network training method based on hybrid semi-supervision and knowledge distillation on a multi-target solid waste dataset.

[0051] This example method randomly samples from the dataset, selects images with labeled data, and uses the remaining unselected parts as unlabeled data images. By comparing the image-based accuracy and object-based accuracy of the grasping detection model in the three training stages, the improvement of model performance in different training stages is reflected.

[0052] In the specific training process, Batch Size is set to 8; the learning rate adopts a dynamic adjustment strategy, which quickly increases the learning rate to 0.005 at the beginning of training, and then gradually decreases to a lower value. The learning rate growth phase accounts for 12%, the initial learning rate scaling factor is 5.5, and the final learning rate scaling factor is 4000. The method of dynamically adjusting the learning rate is used to prevent the learning rate from being set too small or too large, thereby improving the stability and effect of model training; Epoch is set to 50.

[0053] During training, the first stage uses supervised learning methods to train the grasping detection model; the second stage selects the pre-trained model with the best performance in the first stage (with image-based accuracy as the first indicator, and if the image-based accuracy is the same, the object-based accuracy is compared), and hybrid semi-supervised training is performed through random data enhancement; the third stage selects the pre-trained model with the best performance in the second stage as the teacher model, and uses a hybrid semi-supervised method and knowledge distillation method for joint training. The overall flow chart is as follows Figure 7 The accuracy comparison of different stages is shown in Table 3.

[0054] ; The above embodiments are preferred embodiments of the present application. Ordinary technicians in this field can also make various changes or improvements on this basis. Without departing from the overall concept of the present application, these changes or improvements should fall within the scope of protection required by the present application.

Claims

1. A grasp detection training method based on hybrid semi-supervision and knowledge distillation, characterized in that: include: Obtain a training data set, each training data includes an RGB image and a depth feature map of the same captured area; Mark the crawling configuration of some training data in the training data set to build a labeled data set , the remaining training data is used to construct an unlabeled dataset ; For labeled datasets and unlabeled datasets Perform data enhancement to obtain enhanced labeled data sets and enhanced unlabeled datasets; Using labeled datasets Performing supervised training on a first network model pre-built for grasping detection to obtain an initial pre-trained grasping detection model; Use the initial pre-trained grasp detection model to detect the enhanced unlabeled dataset to obtain an enhanced dataset with pseudo labels. ; The enhanced labeled dataset With enhanced pseudo-labeled datasets Mix to get a mixed data set ; Using a mixed dataset , semi-supervised training of the initial pre-trained grasp detection model; The grasp detection model obtained by semi-supervised training is used as the teacher model, and the knowledge distillation method is used based on the mixed dataset. , guiding the training of a second network model used for grasping detection; wherein, the second network model is a lightweight network obtained by simplifying the first network model.

2. The grasping detection training method based on hybrid semi-supervision and knowledge distillation according to claim 1 is characterized in that: The first network model for grasping detection includes an encoder stage, a residual layer stage and a decoder stage; the encoder stage includes a plurality of convolutional layer modules and a plurality of fusion modules respectively following each convolutional layer module, the residual layer stage includes a plurality of residual modules, and the decoder stage includes a plurality of convolutional transposition layers; The second network model for grasping detection is based on the first network model, with the number of fusion modules, residual modules and convolutional transpose layers reduced.

3. The grasping detection training method based on hybrid semi-supervision and knowledge distillation according to claim 2 is characterized in that: The fusion module is an RGB-D fusion module that combines channel attention, spatial attention and channel exchange; in the fusion module: First, the channel attention module processes the RGB image and the depth feature map from the same area and with the same shape to obtain the weight of the channel dimension; Then, the channel exchange module uses the weight information obtained by channel attention to perform channel exchange operations on the RGB image and depth feature map of the input fusion module; The channel swap operation is as follows: if the weight of a feature map in the cth channel is lower than a certain threshold, and the weight of another feature map in the cth channel exceeds another threshold, the cth channel of the former feature map is replaced by the cth channel of the latter feature map; Finally, the spatial attention module processes the feature map after channel swapping to obtain a new color-depth feature map.

4. The grasping detection training method based on hybrid semi-supervision and knowledge distillation according to claim 2 is characterized in that: The residual module is described as: ； ； in, is the input of the residual module, represents convolution, represents regularization, represents global average pooling, and are the feature maps of different intermediate layers, and Both are activation functions; represents element-wise multiplication, is the output of the residual module.

5. The grasping detection training method based on hybrid semi-supervision and knowledge distillation according to claim 1 is characterized in that: Data augmentation performed on the dataset includes random rotation, mirroring, and / or symmetry operations on the images in each training data.

6. The grasping detection training method based on hybrid semi-supervision and knowledge distillation according to claim 1 is characterized in that: Using a mixed dataset , when semi-supervised training is performed on the initial pre-trained grasp detection model, the loss function is expressed as: ； in, Initial pre-trained grasp detection model based on dataset The L1 smooth loss for training, Initial pre-trained grasp detection model based on dataset The L1 smooth loss for training, is the coefficient, is the total loss for semi-supervised training.

7. The grasping detection training method based on hybrid semi-supervision and knowledge distillation according to claim 1 is characterized in that: The loss function for training the second network model is expressed as: ； ； in, is the loss for training the second network model, For the student model based on the mixed dataset The L1 smooth loss between the predicted output and the label; Predict output for the teacher model The predicted output with the student model The L2 loss between them is used to minimize the feature difference between the teacher model and the student model; Input images for the teacher model and the student model, specifically a mixed dataset image training data.

Images