Rock Class Prediction Method Based on Multi-Teacher Knowledge Distillation and Normalized Attention

Abstract

The present invention discloses a rock category prediction method based on multi-teacher knowledge distillation and normalized attention, including: S1: Collect rock sample images, establish a rock data set, and preprocess the original rock sample image data; S2: Insert a normalization-based attention module on the basis of the original network structure to construct a lightweight neural network; S3: Adopt the multi-teacher knowledge distillation method, that is, use the self-attention module represented by Swin-Transformer and the convolutional model represented by ResNet to train a lightweight student model; S4: Divide the data set into a training set and a test set, and record the classification results and various performance indicators. The present invention uses MobileNetV3 as the student model, combines the advantages of the CNN and Self-Attetion paradigms, trains a lightweight model with better performance, obtains the scaling factor through batch normalization, and reflects the change size of each channel through the scaling factor, which also represents the importance of the channel, enabling the implementation of the attention mechanism without additional parameters.

Description

Technical Field

[0001] The present invention belongs to the technical field of image category prediction, and specifically relates to a rock category prediction method based on multi-teacher knowledge distillation and normalized attention. Background Art

[0002] Rocks can be divided into three major categories: igneous rocks, sedimentary rocks, and metamorphic rocks according to their origin. In underground excavation work, the type of rock has an inestimable impact on the progress of the mining work. However, currently in the field of rock category prediction, it is still mainly identified by experienced workers or handed over to professional testing departments for identification. However, in underground excavation work, the efficiency of sending rocks to professional departments for testing is very low, which also poses a great challenge to the smooth progress of the project. In addition, there are the following problems when identifying manually. In actual underground excavation work, it is difficult to collect pictures, and the collected pictures often have a series of problems such as noise and translation, which bring great trouble to rock category prediction.

[0003] For this reason, a method of using a large network for rock category prediction has emerged. Although this method can obtain a relatively good result, due to the poor real-time performance of the large network, and the equipment we use in actual underground excavation work often does not meet the requirements for configuring a large network, so a lightweight network is needed to make its performance meet our needs as much as possible. Summary of the Invention

[0004] The present invention provides a rock category prediction method based on multi-teacher knowledge distillation and normalized attention, and predicts the rock category by combining multi-teacher knowledge distillation and the combination of self-attention and convolution to solve the above technical problems.

[0005] In order to solve the above technical problems, the present invention is implemented as follows:

[0006] The rock category prediction method based on multi-teacher knowledge distillation and normalized attention includes the following steps:

[0007] S1: Collect multiple rock sample images, create a rock data set, and preprocess the original rock sample image data;

[0008] S2: Insert a normalized attention module on the basis of the original network structure to construct a lightweight neural network;

[0009] S3: Adopt the multi-teacher knowledge distillation method, that is, use the self-attention module represented by Swin-Transformer and the convolutional model represented by ResNet to train a lightweight student model;

[0010] S4: Load the rock image to be processed into the model, and record the classification results and various performance metrics.

[0011] The multi-teacher knowledge distillation method combines large neural networks that currently have good effects. The present invention uses a CNN network represented by ResNet and a self-attention network represented by Swin-transformer. The characteristics of these two networks are different: CNN collects information of a region through a square or rectangular receptive field, which is the inductive bias of CNN: ① Local correlation: adjacent regions on the input image will have similar features, and the closer the objects are, the stronger the correlation. ② Translation invariance: CNN can output the same result for the same image and its translated version.

[0012] The Self-Attention model does not introduce additional inductive bias. It uses a weighted average operation based on the context of the input features, and dynamically calculates the attention weights through a similarity function between relevant pixel pairs. And it can obtain global information. This also enables the network to adaptively focus on different regions and capture more features. However, more training data is also required to make the self-attention module achieve the expected effect.

[0013] Therefore, the characteristics of both can be combined to train a neural network model with better performance from the two paradigms.

[0014] For further optimization, in step S1, the preprocessing includes the following steps:

[0015] S11: Classify the rock sample images and place them in the corresponding folders respectively;

[0016] S12: Create two text files, one for the training image paths and their corresponding classes, and the other for the test image paths and their corresponding classes, where the training set accounts for 80% of the dataset and the validation set accounts for 20% of the dataset;

[0017] Step S13: Rewrite the dataset class using Pytorch on jupterlab: Crop the images to a size of 224*224, first take out a batch of image data, calculate the mean and variance of this batch of data, normalize the images according to the mean and variance, and perform data augmentation operations on the training image data with a certain probability; only perform normalization on the test images; then load the data in batches through the DataLoader function and display some image data to observe whether the images can be displayed normally.

[0018] For further optimization, the training image data augmentation operation includes flipping, translating, rotating the training images, and adding random noise. The image augmentation operation is to meet the needs of the actual application environment. When conducting underground excavation, due to a series of reasons such as equipment, lighting, and operation, the rock slices obtained are not as standardized as the data we expect. Therefore, data augmentation can prevent the overfitting of the model.

[0019] For further optimization, in step S2, the construction of the normalized attention module specifically includes the following steps:

[0020] S21: Redesign the attention mechanism submodel based on MobileNetV3, and use the scaling factor γ in batch normalization, as shown in the following formula:

[0021]

[0022] In the formula, μ B and are the mean and variance of the mini-batch samples respectively, γ and β are trainable affine transformation parameters, B in is the input sample, B out is the output sample, ∈ is a small value to prevent division by zero operation;

[0023] S22: Obtain the weight of each channel through the following formula:

[0024]

[0025] In the formula, w i is the weight of the i-th channel, γ i and γ j are the scaling factors of the i-th and j-th channels respectively;

[0026] S23: According to the obtained weights, perform weighted summation with the original feature map, and then obtain the output through a non-linear activation:

[0027] M out = sigmoid(W γ (BN(M in )))

[0028] In the formula, M in is the input feature map, M out is the output feature map, W γ is the weighted summation of the normalized feature map through w i , and sigmoid(·) is the non-linear activation function.

[0029] For further optimization, in step S3, the specific training steps of the multi-teacher knowledge distillation method include:

[0030] S31: Import the image data in the self-built rock slice dataset into the two teacher models after ImageNet pre-training, divide the obtained probability distributions by the temperature parameters in sequence to achieve softening, quantify the confidence of each teacher model through the true labels, and can adaptively weight the predictions through different weights, and obtain the soft labels after softmax processing;

[0031] S32: Import the image data in the self-built rock slice dataset into the student model MobileNetV3-NAM for parallel computing, adopt the same calculation process as the teacher model to obtain soft predictions, and directly input softmax without setting the temperature parameter to obtain hard predictions at the same time;

[0032] S33: Use different weight parameters to obtain the total mixed loss values for the soft labels and soft predictions, hard labels and hard predictions of the obtained multi-teacher models in sequence;

[0033] S34: Use the Adam optimizer and the mixed loss function to train the student model MobileNetV3-NAM, and update the parameters of the student model by means of gradient backpropagation;

[0034] S35: Repeat the above steps until the mixed function and the grading accuracy change curve of the MobileNetV3-NAM model tend to be stable, and end the training of the model.

[0035] For further optimization, the process of obtaining the soft labels is as follows:

[0036]

[0037] In the formula, represents the prediction of the K-th teacher model, and T represents the distillation temperature;

[0038] The process of obtaining the soft predictions is as follows:

[0039]

[0040] The process of obtaining the hard predictions is as follows:

[0041]

[0042] In the formula, represents the prediction of the K-th student model, and T represents the distillation temperature;

[0043] The weighted soft labels are obtained as follows:

[0044]

[0045] The KL divergence between the soft prediction and the soft label is obtained as follows:

[0046]

[0047] The cross - entropy loss function between the hard prediction and the hard label is obtained as follows:

[0048] L S = CrossEntropyLoss(y c , Q SH )

[0049] The total mixed loss value is obtained from the soft labels and soft predictions, hard labels and hard predictions of the multi - teacher model:

[0050] L mix = α·L S +(1 - α)·T 2 ·L T

[0051] Where α represents the weight of the hard label and the hard prediction, and T represents the distillation temperature.

[0052] For further optimization, in step S31, the confidence of each teacher model is quantified by the true label, including the following steps:

[0053] Calculate the cross - entropy between the teacher model prediction and the true label to assign different weights to reflect the weights of different teacher models;

[0054]

[0055]

[0056] Where Q TS，k represents the soft prediction of the K - th teacher model, y represents the true label, L k represents the difference between the prediction of the K - th teacher model and the true label, w k represents the weight of the K - th teacher model, and the smaller L k , the larger w k .

[0057] Compared with the prior art, the beneficial effects of this application are as follows:

[0058] 1. The present invention adopts the method of multi-teacher knowledge distillation, and proposes to use the ResNet and Swin-Transformer models as teacher models, with MobileNetV3 as the student model. By combining the advantages of the CNN and Self-Attetion paradigms, a lightweight model with better performance is trained. Both teacher models need to be pre-trained on the ImageNet dataset and then transferred to the rock dataset for training. A new attention mechanism is introduced. Based on the original channel attention mechanism of MobileNetV3, the scaling factor obtained through batch normalization reflects the change magnitude of each channel and also represents the importance of the channel, enabling the implementation of the attention mechanism without additional parameters.

[0059] 2. Before model training, data augmentation such as rotation and scaling is required, and noise can also be added to the data to simulate rock image data in the actual environment, making the model more robust. Brief Description of the Drawings

[0060] Figure 1 It is a schematic diagram of the basic process of the rock category prediction method based on multi-teacher knowledge distillation and normalization attention according to the present invention;

[0061] Figure 2 It is a schematic diagram of the dataset file according to the present invention;

[0062] Figure 3 It is a schematic diagram of the rock slice image after data augmentation in the present invention;

[0063] Figure 4 It is a schematic diagram of the attention module according to the present invention;

[0064] Figure 5 It is a schematic diagram of the MobileNetV3-NAM model according to the present invention;

[0065] Figure 6 It is a schematic diagram of the training process using multi-teacher knowledge distillation in the present invention;

[0066] Figure 7 It is a schematic diagram of the accuracy rate of 100 rounds of iteration when training MobileNetV3-NAM using the multi-teacher knowledge distillation mode for MobileNetV3-NAM respectively;

[0067] Figure 8 It is a comparative experiment diagram of the accuracy rate of 100 rounds of iteration when training MobileNetV3 using the multi-teacher knowledge distillation mode and training MobileNetV3-NAM using the multi-teacher knowledge distillation mode respectively. Detailed Embodiments

[0068] To make the objectives, technical solutions, and advantages of the embodiments of the present invention clearer, the technical solutions of the present invention will be clearly and completely described below with reference to the accompanying drawings. Apparently, the described embodiments are some, but not all, of the embodiments of the present invention. All other embodiments obtained by those of ordinary skill in the art based on the embodiments of the present invention without creative efforts shall fall within the scope of protection of the present invention.

[0069] A rock category prediction method based on multi-teacher knowledge distillation and normalized attention includes the following steps. The flow of the entire example is as Figure 1 shown:

[0070] Step 1: Collect a rock dataset. The training dataset this time comes from the microscopic image set of rock teaching samples at Nanjing University. Perform certain preprocessing on the original dataset. The results after processing are as Figure 2 shown. The preprocessing includes the following steps: Divide the images into three different categories: igneous rocks, sedimentary rocks, and metamorphic rocks; and place them in three different folders respectively. The number of pictures in each category is: 963, 972, and 699 respectively. Create two text files, one for the training picture paths and their corresponding categories, and the other for the test picture paths and their corresponding categories. Among them, the training set accounts for 80% of the dataset, and the validation set accounts for 20% of the dataset;

[0071] Step 2: Rewrite the dataset class using Pytorch on jupterlab: Crop the pictures to a size of 224*224. First, take out a batch of picture data, calculate the mean and variance of this batch of data, and normalize the pictures according to the mean and variance. Perform a series of data augmentation operations such as flipping, translating, rotating, and adding random noise to the training images with a certain probability. Only perform normalization processing on the test images. Then load the data in batches through the DataLoader function. And display some image data to observe whether the images can be displayed normally, as Figure 3 shown.

[0072] Step 3: First pre-train ResNet-50 and Swin-Transformer-small using ImageNet, and then use the above rock training set to train the data to obtain the weight parameters of the two teacher models. Then build a student model, as Figure 4 shown. Here, it is improved based on MobileNetV3, and a normalized attention module is used to replace the original channel attention module. Then build a complete MobileNetV3-NAM model, as Figure 5 shown. Specifically as follows:

[0073] Use the scaling factor γ in batch normalization, as shown in the following formula.

[0074]

[0075] Among them, μ B and are the mean and variance of the small batch of samples respectively, and γ and β are trainable radiometric transformation parameters. B in is the input sample, B out is the output sample, ∈ is a small value to prevent division by zero operation;

[0076] Then, the weights of each channel are obtained through the following formula:

[0077]

[0078] Among them, w i is the weight of the i-th channel, γ i and γ j are the scaling factors of the i-th and j-th channels respectively;

[0079] Finally, according to the obtained weights, weighted summation is performed with the original feature map, and then through a non-linear activation to obtain the output.

[0080] M out = sigmoid(W γ (BN(M in )))

[0081] Among them, M in is the input feature map, M out is the output feature map, W γ is the weighted summation of the normalized feature map through w i , and sigmoid(·) is a non-linear activation function.

[0082] Step 4: As Figure 6 shown, the preprocessed data is loaded into two teacher models to obtain soft labels. Among them, the process of obtaining soft labels is:

[0083]

[0084] Among them, represents the prediction of the K-th teacher model, and T = 10 represents the distillation temperature.

[0085] The preprocessed data is loaded into the student model to obtain soft predictions and hard predictions. Among them, the process of obtaining soft predictions is:

[0086]

[0087] The process of obtaining hard predictions is:

[0088]

[0089] Wherein: represents the prediction of the K-th student model, and T = 10 represents the distillation temperature.

[0090] The total mixed loss value is obtained by using different weight parameters for the soft labels and soft predictions, hard labels and hard predictions of the obtained multi-teacher model in turn, as follows:[[]]

[0091] Obtain the weighted soft labels as follows:

[0092]

[0093] Obtain the KL divergence between the soft prediction and the soft label as follows:

[0094]

[0095] Obtain the cross-entropy loss function between the hard prediction and the hard label as follows:

[0096] L S = CrossEntropyLoss(y c , Q SH )

[0097] Obtain the total mixed loss value of the soft labels and soft predictions, hard labels and hard predictions of the multi-teacher model:

[0098] L mix = α · L S + (1 - α) · T 2 · L T

[0099] Wherein: α represents the weight of the hard label and the hard prediction.

[0100] Step 5: Use the warmup learning rate, with the warmup steps being 10. The learning rate gradually converges to the minimum value within 100 epochs, and use the Adam optimizer and L mix to train the student model. The initial learning rate is set to 0.0001. With the help of gradient backpropagation, the transfer of the dark knowledge of the multi-teacher model and the update of the parameters of the student model are realized, and the loss value of each epoch is recorded. After each epoch training is completed, a test set is tested once, the test accuracy of each epoch is recorded, and the generalization performance of the test model is evaluated.

[0101] Step 5: Conduct a comparative experiment, and compare the schematic diagrams of the changes in the loss values and the accuracy on the test set with the number of iterations for using knowledge distillation and not using knowledge distillation, as well as MobileNetV4-small-NAM and MobileNetV3-small respectively.

[0102] As Figure 7 shown, two methods were compared: ① Directly using MobileNetV3-NAM for training, after 100 rounds of iteration, the highest accuracy of the obtained student model was 83.77%. ② Using the multi-teacher knowledge distillation mode to train MobileNetV3-NAM, after 100 rounds of iteration, the highest accuracy of the obtained student model was 87.4%. It can be seen that using the multi-teacher knowledge distillation algorithm can effectively improve the accuracy of the model.

[0103] As Figure 8 shown, two methods were compared: ① Using the multi-teacher knowledge distillation mode to train MobileNetV3, after 100 rounds of iteration, the highest accuracy of the obtained student model was 85.9%. Therefore, the classification accuracy can be significantly improved by using the algorithm described in the present invention. ② Using the multi-teacher knowledge distillation mode to train MobileNetV3-NAM, after 100 rounds of iteration, the highest accuracy of the obtained student model was 87.4%. It can be seen that using the normalization-based attention module, namely NAM, is also beneficial to improving the accuracy of the model.

[0104] As shown in Table 1, various parameters of various models were compared, including the number of model parameters: reflecting the total number of internal parameters of the model, used to measure the size of the model; the model calculation amount: reflecting the number of floating-point operations of the model, used to measure the computational complexity of the model; and the MAdds of the model, that is, the multiply-accumulate operation number, reflecting the number of multiplication and addition operations of the model.

[0105] The number of parameters, calculation amount, and MAdds of the optimal model MobileNetV3-NAM described in the present invention are 1.06M, 76.78M, and 151.03M respectively. Compared with the teacher models Resnet-101 and Swin-transformer-Tiny, the complexity of the model is greatly reduced. Therefore, it can be used as a lightweight network and embedded in the equipment used for actual underground excavation work to improve the work efficiency of underground excavation.

[0106] Table 1: Comparison of various parameters of various models described in the present invention

[0107] Model Name Number of Parameters Computational Complexity Multiply-Accumulate Operations MobileNetV3-NAM 1.06M 76.78M 151.03M MobileNetV3 1.52M 76.5M 150.46M Resnet-101 42.51M 10.25G 20.45G Swin-transformer-Tiny 27.52M 7.09G 14.24G

[0108] Taking the ideal embodiments based on the present invention as inspiration, through the above description, relevant staff can completely make various changes and modifications without departing from the technical idea of the present invention. The technical scope of the present invention is not limited to the content in the specification, and its technical scope must be determined according to the scope of the claims.

Claims

1. A rock category prediction method based on multi-teacher knowledge distillation and normalized attention, characterized in that, It includes the following steps: S1: Collect multiple rock sample images, create a rock dataset, and preprocess the original rock sample image data; S2: Insert a normalization-based attention module on the basis of the original network structure to construct a lightweight neural network; specifically, it includes the following steps: S21: Redesign the attention mechanism submodel on the basis of MobileNetV3, and use the scaling factor γ in batch normalization, as shown in the following formula: where μ B and are the mean and variance of the small batch samples respectively, γ and β are trainable radiometric transformation parameters, B in is the input sample, B out is the output sample, ∈ is a small value to prevent division by zero; S22: Obtain the weight of each channel through the following formula: where w i is the weight of the i-th channel, and γ i and γ j are the scaling factors of the i-th and j-th channels respectively; S23: According to the obtained weights, perform weighted summation with the original feature map, and then obtain the output through a non-linear activation: M out = sigmoid(W γ (BN(M in ))) Where, M in is the input feature map, M out is the output feature map, W γ is obtained by weighted summation of the normalized feature map through w i , and sigmoid(·) is the non-linear activation function; S3: Adopt the multi-teacher knowledge distillation method, that is, adopt the self-attention module represented by Swin-Transformer and the convolutional model represented by ResNet to train a lightweight student model; The specific training steps include: S31: Import the image data in the self-built rock slice dataset into two teacher models pre-trained with ImageNet, and divide the obtained probability distributions by the temperature parameter in turn to achieve softening. Quantify the confidence of each teacher model through the true label, and can adaptively weight the prediction through different weights. After softmax processing, obtain the soft label; S32: Import the image data in the self-built rock slice dataset into the student model MobileNetV3-NAM for parallel computing, and adopt the same calculation process as the teacher model to obtain soft predictions, and at the same time directly input softmax without setting the temperature parameter to obtain hard predictions; S33: Use different weight parameters to obtain the total mixed loss value for the soft labels and soft predictions, hard labels and hard predictions of the obtained multi-teacher models in turn; S34: Use the Adam optimizer and the mixed loss function to train the student model MobileNetV3-NAM, and update the parameters of the student model by means of gradient backpropagation; S35: Repeat the above steps until the mixed function and the classification accuracy change curve of the MobileNetV3-NAM model tend to be stable, and end the training of the model; S4: Load the rock image to be processed into the model, and record the classification results and various performance indicators.

2. The rock category prediction method based on multi-teacher knowledge distillation and normalized attention according to claim 1, wherein In the step S1, the preprocessing includes the following steps: S11: Classify the rock sample images and place them in the corresponding folders respectively; S12: Create two text files, one is the training image path and its corresponding category, and the other is the test image path and its corresponding category, where the training set accounts for 80% of the dataset and the validation set accounts for 20% of the dataset; Step S13: Rewrite the dataset class using Pytorch on jupterlab: Crop the pictures to 224*224 size, first take out a batch of picture data, calculate the mean and variance of the batch data, and normalize the pictures according to the mean and variance. Perform data augmentation operations on the training image data with a certain probability; Only perform normalization on the test images; then load the data in batches through the DataLoader function and display some image data to observe whether the images can be displayed normally.

3. The rock category prediction method based on multi-teacher knowledge distillation and normalized attention according to claim 2, wherein The training image data augmentation operations include flipping, translating, rotating, and adding random noise to the training images.

4. The rock category prediction method based on multi-teacher knowledge distillation and normalized attention according to claim 3, wherein The process of obtaining the soft labels is as follows: wherein, represents the prediction of the K-th teacher model, and T represents the distillation temperature; The process of obtaining the soft predictions is as follows: The process of obtaining the hard predictions is as follows: wherein, represents the prediction of the K-th student model, and T represents the distillation temperature; The weighted soft labels are obtained as follows: The KL divergence between the soft predictions and the soft labels is obtained as follows: The cross-entropy loss function between the hard predictions and the hard labels is obtained as follows: L S = CrossEntropyLoss(y c , Q SH ) The total mixed loss value is obtained by obtaining the soft labels, soft predictions, hard labels, and hard predictions of the multi-teacher model: L mix = α·L S + (1 - α)·T 2 ·L T Where α represents the weight of the hard label and the hard prediction, T represents the distillation temperature, and Q TS，k represents the soft prediction of the K-th teacher model, and w k represents the weight of the K-th teacher model.

5. The rock category prediction method based on multi-teacher knowledge distillation and normalized attention according to claim 4, characterized in that, In step S31, the confidence of each teacher model is quantified by the true labels, including the following steps: Calculate the cross-entropy between the teacher model predictions and the true labels to assign different weights to reflect the weights of different teacher models; where Q TS，k represents the soft prediction of the K-th teacher model, y represents the true label, and L k represents the difference between the prediction of the K-th teacher model and the true label, and w k represents the weight of the K-th teacher model. The smaller L k is, the larger w k is.

Images