Image recognition model training method and system, electronic equipment and storage medium

Abstract

The invention provides an image recognition model training method and system, electronic equipment and a storage medium, and the method comprises the steps: obtaining a pre-training data set, a fine tuning data set and a distillation data set, training a preset large pre-training model through the pre-training data set, and obtaining a pre-training weight; initializing a preset fine tuning model based on the pre-training weight, performing fine tuning on the initialized fine tuning model through the fine tuning data set to obtain a teacher model, and performing response distillation on a preset student model through the teacher model and the distillation data set to obtain a target image recognition model; the image recognition model which is efficient, accurate and easy to deploy is provided through the method, and the excellent performance and practicability of the image recognition model can be guaranteed in various actual application environments where computing resources are limited.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of model training, and particularly relates to an image recognition model training method and system, an electronic device and a storage medium. BACKGROUND

[0002] Inspired by natural language large-scale self-supervised pre-training, Masked Image Modeling (MIM) as a self-supervised pre-training method has gradually emerged and become a research hotspot in the field of computer vision. The mainstream training paradigm of this method is to first pre-train on a large amount of unlabeled data, and then fine-tune in the downstream task. In order to improve the generalization ability of the model, a large amount of training data is usually matched with a large model in the pre-training stage, but due to the limitation of computing resources, it is difficult to directly deploy a large visual model to practical applications. In order to solve this problem, the knowledge of the large model is usually transferred to a small model with simple structure and small parameter quantity by means of knowledge distillation, so that the small model reduces the computational amount and storage cost while the performance is close to that of the large model, meeting the needs of practical applications.

[0003] Due to the needs of actual deployment, the training paradigm of most industrial fields currently still relies on supervised learning models pre-trained on open-source general data sets such as ImageNet, and adapts to the field data through fine-tuning. This way cannot fully exert the strong generalization ability of large-scale pre-training, and the performance is limited due to the limitation of model capacity. In addition, although large-scale unsupervised pre-training methods have achieved good results on public data sets (such as ImageNet, COCO, ADE20K, etc.), these data sets have large differences in data distribution with industrial scene data. Therefore, how to effectively connect large-scale unsupervised pre-training and downstream task deployment in the industrial scene has become a key problem to be solved.

[0004] It should be noted that the above content only provides background technical information related to the present application, and does not necessarily constitute prior art. SUMMARY

[0005] The present application provides an image recognition model training method, system, electronic device and storage medium to solve the technical problem of poor generalization ability and performance of image recognition model training.

[0006] Other characteristics and advantages of the present application will become apparent from the following detailed description, or will be learned partly through practice of the present application.

[0007] The application provides an image recognition model training method, including: obtaining a pre-training data set, a fine-tuning data set and a distillation data set; training a pre-set pre-training large model through the pre-training data set to obtain pre-training weights; initializing a pre-set fine-tuning model based on the pre-training weights, and fine-tuning the initialized fine-tuning model through the fine-tuning data set to obtain a teacher model; and performing response distillation on a pre-set student model through the teacher model and the distillation data set to obtain a target image recognition model.

[0008] In an embodiment of the application, based on the foregoing scheme, the pre-training data set, the fine-tuning data set and the distillation data set are obtained by: obtaining first image data of multiple industrial scenes, second image data of general scenes and third image data of preset scenes; mixing the first image data and the second image data to obtain the pre-training data set; and labeling the third image data to obtain the fine-tuning data set and the distillation data set.

[0009] In an embodiment of the application, based on the foregoing scheme, the pre-training weights are obtained by training the pre-set pre-training large model through the pre-training data set, including: training the pre-training large model through a mask reconstruction objective function and the pre-training data set to obtain the pre-training weights, wherein the mask reconstruction objective function is determined based on real pixel values and predicted pixel values of an input image.

[0010] In an embodiment of the application, based on the foregoing scheme, the teacher model is obtained by initializing the pre-set fine-tuning model based on the pre-training weights and fine-tuning the initialized fine-tuning model through the fine-tuning data set, including: initializing an encoder of the fine-tuning model based on the pre-training weights, and randomly initializing a decoder of the fine-tuning model to obtain the initialized fine-tuning model; training the initialized fine-tuning model through a first target loss function and the fine-tuning data set, and using a layer-by-layer decreasing learning rate adjustment strategy to train the encoder of the initialized fine-tuning model to obtain the teacher model, wherein the first target loss function is a combination function of cross-entropy loss and intersection over union loss.

[0011] In an embodiment of the application, based on the foregoing scheme, the method further includes: obtaining an initial feature map of a target preset multiple from an original image through the encoder, wherein the encoder sequentially includes a first stage layer of a first preset multiple down-sampling from top to bottom, and a second stage layer of a preset number of second preset multiple down-sampling; and obtaining a target feature map of an original image size from the initial feature map through the decoder, wherein the decoder sequentially includes an up-sampling layer, a skip link layer and a convolution layer from bottom to top.

[0012] In an embodiment of the present application, based on the foregoing scheme, the preset student model is responsively distilled by the teacher model and the distillation dataset to obtain a target image recognition model, including: inputting the distillation dataset into the teacher model and the student model respectively to obtain a first output probability distribution of the teacher model and a second output probability distribution of the student model respectively; determining the relative entropy of the first output probability distribution and the second output probability distribution, and determining a distillation target function based on the relative entropy; responsively distilling the student model based on a second target loss function to obtain the target image recognition model, the second target loss function being a combination function of cross-entropy loss, intersection over union loss and the distillation target function.

[0013] In an embodiment of the present application, based on the foregoing scheme, the method further includes: obtaining test data; pre-processing the test data to obtain pre-processed test data; inputting the pre-processed test data into the target image recognition model for image recognition to obtain a recognition result; and converting the recognition result into a binary image and storing it.

[0014] The present application also provides an image recognition model training system, including: a collection module configured to obtain a pre-training dataset, a fine-tuning dataset and a distillation dataset; a training module configured to train a preset pre-training large model by the pre-training dataset to obtain pre-training weights; a fine-tuning module configured to initialize a preset fine-tuning model based on the pre-training weights, and fine-tune the initialized fine-tuning model by the fine-tuning dataset to obtain a teacher model; and a distillation module configured to responsively distill a preset student model by the teacher model and the distillation dataset to obtain a target image recognition model.

[0015] The present application also provides an electronic device including a memory and a processor, the memory storing a computer program, and the processor implementing the image recognition model training method of any one of the above embodiments when executing the computer program.

[0016] The present application also provides a computer-readable storage medium storing a computer program, and the computer program is executed by a processor to implement the image recognition model training method of any one of the above embodiments.

[0017] Beneficial effects of the present application: by introducing a new paradigm of large-scale self-supervised pre-training, fine-tuning and distillation, the generalization ability of image recognition models in industrial scenarios can be effectively improved, while solving the problem of large model deployment difficulty, by fusing the first image data of the industrial scenario and the second image data of the general scenario in the pre-training stage, using the mask image modeling method for self-supervised pre-training, the pre-training weight with good generalization is provided for the downstream task; in the fine-tuning stage, combined with the layer-by-layer decreasing learning rate adjustment strategy, the performance of the model on the industrial scene data is further improved; finally, through knowledge distillation, the knowledge of the teacher model is transferred to the student model, so that the student model can maintain high recognition accuracy while having small computational complexity and storage overhead, and can be successfully deployed in the industrial scenario.

[0018] It should be understood that the above general description and the following detailed description are only exemplary and explanatory, and cannot limit the present application. BRIEF DESCRIPTION OF DRAWINGS

[0019] The drawings incorporated into the specification and forming part of the specification, show embodiments consistent with the present application, and together with the specification, serve to explain the principles of the present application. It is obvious that the drawings in the following description are only some embodiments of the present application, and those skilled in the art can obtain other drawings from these drawings without creative labor.

[0020] In the drawings: Figure 1 is a schematic diagram of an exemplary system architecture shown by an exemplary embodiment of the present application; Figure 2 is a flowchart of an image recognition model training method shown by an exemplary embodiment of the present application; Figure 3 is a flowchart of an image recognition model training method shown by another exemplary embodiment of the present application; Figure 4 is a framework diagram of an image recognition model training method shown by an exemplary embodiment of the present application; Figure 5 is a block diagram of an image recognition model training system shown by an exemplary embodiment of the present application; Figure 6 shows a structural schematic diagram of a computer system of an electronic device suitable for implementing embodiments of the present application. DETAILED DESCRIPTION

[0021] The following detailed description discusses embodiments of the application in connection with specific specific examples. Those skilled in the art will readily recognize that other advantages and benefits can be realized from the techniques described herein. The application can be implemented in any number of manners, including the specific embodiments described below, as well as in other ways not specifically described herein. The detailed description is presented primarily for the purpose of illustrating the principles and advantages of the application and not for limiting the same.

[0022] It should be noted that the figures provided in the following embodiments are only schematically illustrating the basic concept of the application, and the figures only show the components related to the application in the drawings, not the number, shape and size of the components when actually implemented. The actual implementation of each component may be a random change in type, number and proportion, and the layout type of the components may also be more complex.

[0023] In the following description, a large number of details are discussed to provide a more thorough explanation of the embodiments of the application. However, it is obvious to those skilled in the art that the embodiments of the application can be implemented without these specific details, and in other embodiments, well-known structures and devices are shown in block diagrams rather than in detail to avoid obscuring the embodiments of the application.

[0024] First of all, it should be noted that IOU (Intersection over Union, Intersection over Union loss) is an evaluation index and loss function commonly used in object detection and semantic segmentation tasks. It measures the degree of overlap between the predicted bounding box (or mask) and the true bounding box (or mask). IOU is defined as the ratio of the intersection to the union. Specifically, given two bounding boxes A and B, IOU is the intersection area of A and B divided by their union area. Its value ranges from 0 to 1, and the larger the value, the higher the degree of overlap between the two bounding boxes.

[0025] Relative entropy, also known as Kullback-Leibler divergence or information divergence, is an asymmetric measure of the difference between two probability distributions. In information theory, relative entropy is equivalent to the difference in Shannon entropy between two probability distributions. Relative entropy is the loss function for optimization algorithms such as the Expectation-Maximization algorithm. In this case, one probability distribution is the true distribution, and the other is a theoretical (fitted) distribution. Relative entropy represents the information loss that occurs when fitting the true distribution with the theoretical distribution. KL divergence is a fundamental measure of probability distribution distance, commonly used in machine learning, statistical inference, and various scenarios in data science.

[0026] A binary image is an image in which each pixel has only two possible values or gray level states. Binary images are often represented as black and white, B&W, or monochrome images. In a binary image, there are only two gray levels, meaning that any pixel in the image has a gray value of 0 or 255, representing black and white, respectively.

[0027] SegModel is a deep learning model for image semantic segmentation, mainly applied in medical image analysis, autonomous driving, and other fields. Its core network structure is U-Net, which realizes pixel-level classification through an encoder-decoder structure.

[0028] Figure 1 is a schematic diagram of an exemplary system architecture shown in an exemplary embodiment of the present application.

[0029] Referring to Figure 1As shown, the system architecture can include a data acquisition device 101 and a computer device 102. Among them, the computer device 102 can be at least one of a desktop graphic processing unit (GPU) computer, a GPU computing cluster, a neural network computer, etc. The data acquisition device 101 is used to collect a pre-training data set, a fine-tuning data set and a distillation data set. In this embodiment, the data acquisition device 101 provides the above data to the computer device 102 for processing after obtaining the data. The computer device 102 can be used by a person skilled in the art to train a pre-set pre-training large model through a pre-training data set to obtain pre-training weights, initialize a pre-set fine-tuning model based on the pre-training weights, fine-tune the initialized fine-tuning model through a fine-tuning data set, obtain a teacher model, and respond to distillation of a pre-set student model through the teacher model and a distillation data set to obtain a target image recognition model, so as to complete image recognition model training. It should be noted that the data acquisition device 101 and the computer device 102 provided in this embodiment are only an example and should not limit the functions and use range of the embodiments of the present application.

[0030] It should be noted that the image recognition model training method provided in the embodiments of the present application is generally executed by the computer device 102, and correspondingly, the image recognition model training system is generally provided in the computer device 102.

[0031] Figure 2 is a flowchart of an image recognition model training method provided in an exemplary embodiment of the present application. The image recognition model training method can be executed by a computing processing device, which can be a computer device 102 as shown in Figure 1 . Referring to Figure 2 , the image recognition model training method at least includes steps S210 to S240, which are described in detail as follows: In step S210, a pre-training data set, a fine-tuning data set and a distillation data set are obtained.

[0032] In an embodiment of the present application, first image data of multiple industrial scenes, second image data of general scenes, and third image data of a pre-set scene are obtained; the first image data and the second image data are mixed to obtain a pre-training data set; the third image data is labeled to obtain a fine-tuning data set and a distillation data set.

[0033] In the embodiment, first image data of various industrial scenes with diversity is collected, and second image data of a general scene is mixed as a preset pre-training data set of a pre-trained large model. In the fine-tuning and distillation stages, third image data of a specific scene, i.e., a production scene, is collected, and unified labeling is performed with the pre-trained large model as an application scene of laser line segmentation, to obtain a fine-tuning data set and a distillation data set for fine-tuning and distillation.

[0034] In step S220, the pre-trained large model is trained by the pre-training data set, to obtain pre-training weights.

[0035] In an embodiment of the application, the pre-training large model is trained by a mask reconstruction objective function and the pre-training data set, to obtain the pre-training weights, wherein the mask reconstruction objective function is determined based on real pixel values and predicted pixel values of an input image.

[0036] In the embodiment, a network with an encoder-decoder architecture of an order of hundreds of millions of parameters is used as the pre-trained large model, large-scale self-supervised pre-training is performed on the pre-training data set in a masked image modeling (MIM) manner, and pre-training weights with certain generalization capability for industrial scenes are obtained. The mask reconstruction objective function of the above-mentioned masked image modeling is generally as follows: The formula is as follows:

[0037] In the formula, x and y are RGB values of an input image, i.e., real pixel values and predicted pixel values, M represents a set of masked pixels, , x, y are RGB values of an input image, i.e., real pixel values and predicted pixel values, M represents a set of masked pixels, is the number of elements, and L is the mask reconstruction objective function. By default, an L1 loss is used.

[0038] The mask reconstruction objective is to measure the difference between the mask reconstruction image and the original image using an L1 loss, and to minimize the difference between the reconstruction image and the original image.

[0039] In step S230, the preset fine-tuning model is initialized based on the pre-training weights, and the initialized fine-tuning model is fine-tuned by the fine-tuning data set, to obtain a teacher model.

[0040] In an embodiment of the present application, the encoder of the fine-tuning model is initialized based on the pre-training weights, the decoder of the fine-tuning model is randomly initialized, to obtain an initialized fine-tuning model; the initialized fine-tuning model is trained through a first target loss function and a fine-tuning data set, and the encoder of the initialized fine-tuning model is trained using a layer-by-layer decreasing learning rate adjustment strategy, to obtain a teacher model, wherein the first target loss function is a combination function of a cross-entropy loss and an intersection-over-union loss.

[0041] In the present embodiment, the semantic segmentation network SegModel is used as the fine-tuning model, the encoder part is defined as Encoder, the decoder part is defined as Decoder, the encoder part of the model is initialized using pre-training weights, the decoder part is randomly initialized, full-scale fine-tuning is performed on the fine-tuning data set, the training loss function, i.e., the first target loss function, is a combination function of a cross-entropy loss and an IoU loss (i.e., an intersection-over-union loss), and a layer-by-layer decreasing learning rate adjustment strategy is used for the encoder part during training to improve fine-tuning precision. It can be understood that the layer-by-layer decreasing learning rate adjustment strategy is used for the encoder part, i.e., different learning rates are used for different layers of the encoder during training (fine-tuning), and the learning rate gradually decreases as the network level deepens.

[0042] In the present embodiment, the weights of the cross-entropy loss and the IoU loss in the first target loss function can be set according to requirements, for example, the cross-entropy loss and the IoU loss weights are set to 1 and 3, respectively.

[0043] In an embodiment of the present application, the encoder is used to obtain an initial feature map of a target preset multiple based on an original image, wherein the encoder sequentially comprises, from top to bottom, a first stage layer of a first preset multiple down-sampling and a preset number of second stage layers of a second preset multiple down-sampling; and the decoder is used to obtain a target feature map of the original image size based on the initial feature map, wherein the decoder sequentially comprises, from bottom to top, an up-sampling layer, a skip connection layer and a convolution layer.

[0044] In the embodiment, the specific structure of the encoder is Stem layer-downsampling by 4-Stage1 layer-Stage2 layer-Stage3 layer-Stage4 layer from top to bottom, each Stage layer contains several Conv Blocks, and each will downsample the image by 2, and finally obtain a feature map downsampled by 32. The number of blocks contained in each Stage and the channel dimension output by each Stgae are (3, 3, 27, 3) and (128, 256, 512, 1024) respectively. It can be understood that the Conv Block is a basic module for building a convolutional neural network (CNN) in deep learning, mainly used for extracting image features; from top to bottom refers to the flow direction of data in the neural network, that is, the process from the input layer to the output layer; the specific structure of the above encoder and the final downsampled multiple are for illustration only, and can be adjusted according to the needs in actual application, and the comparison is not limited by the present application.

[0045] In the embodiment, the specific structure of the decoder is upsampling layer-jump link layer-convolutional layer from bottom to top, and the structure is repeated 4 times to obtain a feature map of the original image size, the jump link layer connects the features transmitted from the encoder, and splices in the channel dimension to fuse multi-scale features and improve the segmentation accuracy. It can be understood that the decoder from bottom to top refers to the flow direction of data in the decoder, and can also be understood as the process of reconstructing and restoring data in the decoder from deep abstract features to shallow specific pixel space.

[0046] In step S240, the pre-set student model is response distilled through the teacher model and the distillation data set to obtain the target image recognition model.

[0047] In an embodiment of the present application, the distillation data set is input into the teacher model and the student model respectively to obtain a first output probability distribution of the teacher model and a second output probability distribution of the student model respectively; the relative entropy of the first output probability distribution and the second output probability distribution is determined, and the distillation target function is determined based on the relative entropy; the student model is response distilled based on the second target loss function to obtain the target image recognition model, and the second target loss function is a combination function of cross-entropy loss, intersection over union loss and distillation target function.

[0048] In the embodiment, the distillation data set is obtained by using the same processing method as the fine-tuning data set, the fine-tuned billion-level parameter model is used as the teacher model, and the SegModel is used as the student model for distillation. A plurality of same pictures are input into the teacher model and the student model respectively to obtain a first output distribution of the teacher model and a second output probability distribution of the student model Then, distillation is performed using a response distillation mode, and the response distillation mode aims to make the outputs of the student model and the teacher model as close as possible to achieve a balance between the parameter quantity and the model performance.

[0049] The distillation mode is specifically to use Kullback Leibler divergence to measure the similarity of the first output probability distribution and the second output probability distribution, to reduce the distance between the first output probability distribution and the second output probability distribution as much as possible, and the distillation target function is: The formula is

[0050] In the formula, is the KL divergence, i.e., the relative entropy, is the first output probability distribution of the teacher model and the second output probability distribution of the student model, respectively, L is the distillation target function, and H and W are the length and width of the output feature map.

[0051] In the distillation process, the student model is not only supervised by the teacher model, but also supervised by the real label. The training loss function, i.e., the second target loss function, is a combination function of the cross-entropy loss, the IoU loss, and the KL divergence loss (i.e., the distillation target function).

[0052] In this embodiment, the weights of the cross-entropy loss, the IoU loss, and the KL divergence loss of the second target loss function can be set according to requirements, for example, the cross-entropy loss, the IoU loss, and the KL divergence loss weights are set to 1, 3, and 0.5, respectively.

[0053] In an embodiment of the present application, test data is obtained; the test data is preprocessed to obtain preprocessed test data; the preprocessed test data is input into a target image recognition model for image recognition to obtain a recognition result; and the recognition result is converted into a binary image and stored.

[0054] In this embodiment, the obtained distillation model, i.e., the target image recognition model, is exported in ONNX format for deployment testing, the test data is preprocessed and input into the above-mentioned model to obtain a recognition result, and the result is stored in the form of a binary image, in which the foreground region pixel value is 1 and the pixel value of other positions is 0. It can be understood that ONNX (Open Neural Network Exchange) is an open deep learning model exchange format, which aims to solve the compatibility problem of model migration and deployment between different frameworks. In the actual application process, the target image recognition model can be exported in other formats according to requirements, and the present application does not limit this.

[0055] Figure 3is a flowchart of an image recognition model training method according to another example embodiment of the present application, referring to Figure 3 As shown in the figure, in an example embodiment, industrial scene data is collected to form a pre-training data set, a fine-tuning data set, and a distillation data set, respectively; in the pre-training stage, the encoder of the pre-training large model is pre-trained through the mask image modeling method and the pre-training data set to obtain pre-training weights; in the fine-tuning stage, the encoder of the fine-tuning model is initialized using the pre-training weights, and the initialized fine-tuning model is fine-tuned through the downstream task using the fine-tuning data set, and the fine-tuned model is used as the teacher model; in the distillation stage, the student model is distilled based on the teacher model and the distillation data set to obtain a lightweight distilled model, i.e., the target image recognition model; and the distilled model is converted into an ONNX format for deployment testing. The pre-training large model is a visual large model. The image recognition model training method in this embodiment belongs to the same concept as the above-mentioned embodiments, and the specific way of performing operations in each step has been described in detail in the foregoing embodiments, which will not be described here.

[0056] Figure 4 is a framework diagram of an image recognition model training method according to an example embodiment of the present application, referring to Figure 4 As shown in the figure, in an example embodiment, in the pre-training stage, the large-parameter visual encoder is pre-trained using the pre-training data set to obtain pre-training weights; in the fine-tuning stage, the encoder is initialized using the pre-training weights, and the initialized fine-tuning model is fine-tuned using the downstream segmentation model and the fine-tuning data set to obtain a teacher model; in the distillation stage, the student model is distilled based on the teacher model and the distillation data set to obtain a lightweight model, i.e., the target image recognition model, which is deployed. The image recognition model training method in this embodiment belongs to the same concept as the foregoing embodiments, and the specific way of performing operations in each step has been described in detail in the foregoing embodiments, which will not be described here.

[0057] In an embodiment of the present application, an embodiment 1 is provided, which continues to refer to Figure 3 and Figure 4 , and is as follows: Step 1, for the pre-training stage, collect first picture data of industrial scenes with diversity, and mix second image data of general scenes as the pre-training data set of the data large model (the mixing ratio can be 6:4). For the fine-tuning and distillation stage, use a total of 20,000 laser line pictures as the training set, 2,000 laser line pictures as the test set, and perform unified labeling as the fine-tuning data set and the distillation data set used in the fine-tuning and distillation of the present application.

[0058] Step 2, using a network of an encoder-decoder architecture with a large number of parameters as a pre-training large model, large-scale self-supervised pre-training on the pre-training data set in step 1 in a way of modeling a masked image, to obtain pre-training weights with certain generalization ability for industrial scenes, and the specific loss function is wherein x and y are the real pixel value and the predicted pixel value of the input image respectively, M represents the set of masked pixels, is the number of elements, and N is the number of pictures per iteration, is the hyperparameter of the loss function.

[0059] Step 3, using a semantic segmentation network SegModel as a fine-tuning model, defining the encoder part as Encoder and the decoder part as Decoder, initializing the encoder part of the model with the pre-training weights obtained in step 2, and randomly initializing the decoder part, performing full-scale fine-tuning on the fine-tuning data set in step 1, and the training loss function is a combination function of cross-entropy loss and IoU loss, specifically wherein, is the cross-entropy loss, is the IoU loss, and T is the hyperparameter of the loss function. In this example, the graphics card used is Nvidia GTX 3090Ti with 24GB of video memory, the data size of one training batch is 8, the initial learning rate is 0.0001, the warmup is 1000 iterations, the learning rate decay rate of each layer of the encoder is 0.9, and a total of 40k iterations are trained to obtain pre-training weights with certain generalization ability for industrial scenes.

[0060] Step 4, using the same data set as step 3 as a distillation data set, using the fine-tuned large model with a large number of parameters in step 3 as a teacher model, and using SegModel as a student model for distillation, inputting the same picture into the teacher model and the student model to obtain respectively, and then distilling them in a response to distillation manner. The training loss function is a combination function of cross-entropy loss, IoU loss and KL divergence loss, specifically wherein, is the cross-entropy loss, is the IoU loss, is the KL divergence loss, and S is the hyperparameter of the loss function. In this example, the graphics card used is Nvidia GTX 3090Ti with 24GB of video memory, the teacher model uses the fine-tuned weights in step 3, the training batch size is uniformly set to 4, the initial learning rate is 0.0001, the warmup is 1000 iterations, and a total of 40k iterations are trained.

[0061] It can be understood that the present embodiment is only for illustration, and the number of pictures, graphics card models, iteration times and the like involved in the above description are not limited by the present application, nor should they bring any limitation to the function and use range of the embodiments of the present application.

[0062] In one embodiment of the present application, an embodiment 2 is provided, which tests the data test set using the model trained in embodiment 1. In order to verify the effectiveness of the present application, the embodiment also compares the SegModel model trained using supervised pre-training weight and the same model as the distillation model trained from scratch. Table 1 is a comparison of the foreground IoU, background IoU, mIoU and model parameter size of several models, and three different SegModel correspond to different parameter size versions. It can be understood that the present embodiment is only for illustration, and the model types and comparison parameters in table 1 can be increased, reduced or changed according to actual needs, and the values can also be modified according to actual needs, which are not limited by the present application. It can be seen that compared with supervised pre-training weight, using large-scale self-supervised pre-training has higher precision, which verifies the effectiveness of self-supervised pre-training in industrial scenarios. At the same time, comparing the models distilled from different groups and the models trained from scratch, it can be seen that the models distilled have higher precision under the same parameter amount. Compared with the teacher model, the distilled model has a smaller decrease in precision while achieving a significant decrease in parameter amount.

[0063] Table 1 Comparison of foreground IoU, background IoU, mIoU and model parameter amount of different models

[0064] It can be understood that the present application still has accurate recognition results while significantly reducing the parameter amount, which verifies the effectiveness of the present application. From top to bottom are the teacher model (large model after fine-tuning), the distilled SegModel model, the SegModel model trained from scratch, the distilled SegModel2 model, the distilled SegModel3 model, the original image and the labeled image.

[0065] The present application effectively solves the problems of computational resource limitation and deployment difficulty of image recognition models in industrial scenarios by combining large-scale self-supervised pre-training, fine-tuning, and knowledge distillation techniques. First, by combining the first image data of the industrial scenario and the second image data of the general scenario in the pre-training stage, a large-scale self-supervised pre-training is performed using a mask image modeling method to obtain a visual large model with good generalization ability. Then, a layer-by-layer decreasing learning rate adjustment strategy is used to fine-tune the model to adapt to the image data in a specific industrial scenario, further improving the model performance. Finally, the knowledge of the large model, i.e., the teacher model, is transferred to the small model, i.e., the student model, through the knowledge distillation method, ensuring that the small model can significantly reduce the computational load and storage overhead while maintaining high recognition accuracy. This method not only improves the accuracy of image recognition tasks in industrial scenarios, but also effectively solves the problem of deploying large models in actual environments with limited computational resources, enabling small models to run efficiently in industrial applications with higher deployment flexibility and adaptability. In summary, the present application provides an efficient, accurate, and easy-to-deploy image recognition model training method for the industrial field, which can ensure the excellent performance and practicality of image recognition models in various actual application environments with limited computational resources.

[0066] Figure 5 is a block diagram of an image recognition model training system according to an example embodiment of the present application. The system can be applied to Figure 1 the implementation environment shown in FIG. 1, and is specifically configured in the computer device 102. The system can also be applied to other example implementation environments and is specifically configured in other devices, and the implementation environment to which the system is applied is not limited by the present embodiment.

[0067] As shown in Figure 5 , the example image recognition model training system includes a collection module 510, a training module 520, a fine-tuning module 530, and a distillation module 540.

[0068] The collection module 510 is configured to obtain a pre-training data set, a fine-tuning data set, and a distillation data set. The training module 520 is configured to train a pre-set pre-training large model using the pre-training data set to obtain pre-training weights. The fine-tuning module 530 is configured to initialize a pre-set fine-tuning model based on the pre-training weights, and fine-tune the initialized fine-tuning model using the fine-tuning data set to obtain a teacher model. The distillation module 540 is configured to perform response distillation on a pre-set student model using the teacher model and the distillation data set to obtain a target image recognition model.

[0069] It should be noted that the image recognition model training system provided in the above embodiments and the image recognition model training method provided in the above embodiments belong to the same concept, wherein the specific manner in which each module and unit performs operations has been described in detail in the method embodiments, which will not be described here. The image recognition model training system provided in the above embodiments can be used in actual applications, and the above functions can be completed by different functional modules according to needs, that is, the internal structure of the system is divided into different functional modules to complete all or part of the functions described above, and this is not limited herein.

[0070] Embodiments of the present application also provide an electronic device comprising a memory and a processor, wherein the memory stores a computer program, and the processor implements the image recognition model training method according to any one of the above embodiments when executing the computer program.

[0071] Figure 6 The structure of a computer system of an electronic device suitable for implementing embodiments of the present application is shown. It should be noted that, Figure 6 The computer system 600 of the electronic device shown is only an example and should not limit the functions and use range of the embodiments of the present application.

[0072] As Figure 6 shown, the computer system 600 includes a central processing unit (CPU) 601, which can perform various appropriate actions and processes according to programs stored in a read-only memory (ROM) 602 or programs loaded from a storage portion 608 to a random access memory (RAM) 603, such as the methods provided in the above embodiments. In the RAM 603, various programs and data required for system operation are also stored. The CPU 601, the ROM 602, and the RAM 603 are connected to each other through a bus 604. An input / output (I / O) interface 605 is also connected to the bus 604.

[0073] The following components are connected to I / O interface 605: an input section 606 including a keyboard, mouse, etc.; an output section 607 including a cathode ray tube (CRT), liquid crystal display (LCD), etc., and speakers, etc.; a storage section 608 including a hard disk, etc.; and a communication section 609 including a network interface card such as a LAN (Local Area Network) card, modem, etc. The communication section 609 performs communication processing via a network such as the Internet. A drive 610 is also connected to I / O interface 605 as needed. A removable medium 611, such as a disk, optical disk, magneto-optical disk, semiconductor memory, etc., is installed on drive 610 as needed so that computer programs read from it can be installed into storage section 608 as needed.

[0074] Specifically, according to embodiments of this application, the processes described above with reference to the flowcharts can be implemented as computer software programs. For example, embodiments of this application include a computer program product comprising a computer program carried on a computer-readable medium, the computer program including a computer program for performing the methods shown in the flowcharts. In such embodiments, the computer program can be downloaded and installed from a network via communication section 609, and / or installed from removable medium 611. When the computer program is executed by central processing unit (CPU) 601, it performs various functions defined in the system of this application.

[0075] It should be noted that the computer-readable medium shown in the embodiments of this application can be a computer-readable signal medium or a computer-readable storage medium, or any combination of the two. A computer-readable storage medium can be, for example, an electrical, magnetic, optical, electromagnetic, infrared, or semiconductor system, apparatus, or device, or any combination thereof. More specific examples of a computer-readable storage medium may include, but are not limited to: an electrical connection having one or more wires, a portable computer disk, a hard disk, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM), flash memory, optical fiber, portable compact disc read-only memory (CD-ROM), optical storage device, magnetic storage device, or any suitable combination thereof. In this application, a computer-readable signal medium may include a data signal propagated in baseband or as part of a carrier wave, carrying a computer-readable computer program. Such propagated data signals can take various forms, including but not limited to electromagnetic signals, optical signals, or any suitable combination thereof. Computer-readable signal media can also be any computer-readable medium other than computer-readable storage media, which can send, propagate, or transmit a program for use by or in connection with an instruction execution system, apparatus, or device. The computer program contained on the computer-readable medium can be transmitted using any suitable medium, including but not limited to wireless, wired, etc., or any suitable combination thereof.

[0076] The flowcharts and block diagrams in the accompanying drawings illustrate the architecture, functionality, and operation of possible implementations of systems, methods, and computer program products according to various embodiments of this application. Each block in a flowchart or block diagram may represent a module, segment, or portion of code, which contains one or more executable instructions for implementing a specified logical function. It should also be noted that in some alternative implementations, the functions indicated in the blocks may occur in a different order than those indicated in the drawings. For example, two consecutively indicated blocks may actually be executed substantially in parallel, and they may sometimes be executed in reverse order, depending on the functions involved. It should also be noted that each block in a block diagram or flowchart, and combinations of blocks in a block diagram or flowchart, can be implemented using a dedicated hardware-based system that performs the specified function or operation, or using a combination of dedicated hardware and computer instructions.

[0077] The units described in the embodiments of this application can be implemented in software or hardware, and the described units can also be located in a processor. The names of these units do not necessarily limit the specific unit itself.

[0078] Another aspect of this application provides a computer-readable storage medium storing a computer program that is executed by a processor to implement the image recognition model training method as described in any of the above embodiments. This computer-readable storage medium may be included in the electronic device described in the above embodiments, or it may exist independently and not assembled into the electronic device.

[0079] It should be noted that although several modules or units for the device used to perform actions have been mentioned in the detailed description above, this division is not mandatory. In fact, according to the embodiments of this application, the features and functions of two or more modules or units described above can be embodied in one module or unit. Conversely, the features and functions of one module or unit described above can be further divided and embodied by multiple modules or units.

[0080] Another aspect of this application provides a computer program product or computer program including computer instructions stored in a computer-readable storage medium. A processor of a computer device reads the computer instructions from the computer-readable storage medium and executes the computer instructions, causing the computer device to perform the image recognition model training method provided in the various embodiments described above.

[0081] Through the above description of the embodiments, those skilled in the art will readily understand that the exemplary embodiments described herein can be implemented by software or by combining software with necessary hardware. Therefore, the technical solutions according to the embodiments of this application can be embodied in the form of a software product, which can be stored in a non-volatile storage medium (such as a CD-ROM, USB flash drive, external hard drive, etc.) or on a network, including several instructions to cause a computing device (such as a personal computer, server, touch terminal, or network device, etc.) to execute the method according to the embodiments of this application.

[0082] Other embodiments of this application will readily occur to those skilled in the art upon consideration of the specification and practice of the embodiments disclosed herein. This application is intended to cover any variations, uses, or adaptations of this application that follow the general principles of this application and include common knowledge or customary techniques in the art not disclosed herein.

[0083] The above embodiments are merely illustrative of the principles and effects of this application and are not intended to limit this application. Any person skilled in the art can modify or alter the above embodiments without departing from the spirit and scope of this application. Therefore, all equivalent modifications or alterations made by those skilled in the art without departing from the spirit and technical concept disclosed in this application should still be covered by the claims of this application.

Claims

1. A method for training an image recognition model, characterized in that, include: Obtain the pre-trained dataset, the fine-tuned dataset, and the distillation dataset; The pre-trained large model is trained using the pre-trained dataset to obtain pre-trained weights. The pre-trained weights are used to initialize the preset fine-tuning model, and the initialized fine-tuning model is then fine-tuned using the fine-tuning dataset to obtain the teacher model. The target image recognition model is obtained by responding to the preset student model using the teacher model and the distillation dataset.

2. The image recognition model training method according to claim 1, characterized in that, Obtain the pre-trained dataset, fine-tuned dataset, and distillation dataset, including: Acquire first image data of various industrial scenarios, second image data of general scenarios, and third image data of preset scenarios; The first image data and the second image data are mixed to obtain the pre-training dataset; The third image data is labeled to obtain the fine-tuning dataset and the distillation dataset.

3. The image recognition model training method according to claim 1, characterized in that, The pre-trained large model is trained using the pre-trained dataset to obtain pre-trained weights, including: The pre-trained large model is trained using a mask reconstruction objective function and the pre-trained dataset to obtain the pre-trained weights, wherein the mask reconstruction objective function is determined based on the true pixel values ​​and predicted pixel values ​​of the input image.

4. The image recognition model training method according to any one of claims 1 to 3, characterized in that, The pre-trained weights are used to initialize a preset fine-tuning model, and the initialized fine-tuning model is then fine-tuned using the fine-tuning dataset to obtain a teacher model, including: The encoder of the fine-tuned model is initialized based on the pre-trained weights, and the decoder of the fine-tuned model is randomly initialized to obtain the initialized fine-tuned model. The initialized fine-tuned model is trained using the first objective loss function and the fine-tuned dataset, and the encoder of the initialized fine-tuned model is trained using a layer-by-layer decreasing learning rate adjustment strategy to obtain the teacher model. The first objective loss function is a combination function of cross-entropy loss and cross-union ratio loss.

5. The image recognition model training method according to claim 4, characterized in that, The method further includes: The encoder obtains an initial feature map of a target preset multiple based on the original image. The encoder consists of a first-stage layer with a first preset multiple downsampling and a second-stage layer with a preset number of second preset multiple downsamplings from top to bottom. The decoder obtains a target feature map of the original image size based on the initial feature map, wherein the decoder consists of an upsampling layer, a skip connection layer, and a convolutional layer from bottom to top.

6. The image recognition model training method according to any one of claims 1 to 3, characterized in that, By performing response distillation on a pre-defined student model using the teacher model and the distillation dataset, a target image recognition model is obtained, including: The distillation dataset is input into the teacher model and the student model respectively, and the first output probability distribution of the teacher model and the second output probability distribution of the student model are obtained respectively. Determine the relative entropy of the first output probability distribution and the second output probability distribution, and determine the distillation objective function based on the relative entropy; The student model is distilled based on the second objective loss function to obtain the target image recognition model. The second objective loss function is a combination of cross-entropy loss, cross-union loss and distillation objective function.

7. The image recognition model training method according to any one of claims 1 to 3, characterized in that, The method further includes: Obtain test data; The test data is preprocessed to obtain preprocessed test data; The preprocessed test data is input into the target image recognition model for image recognition to obtain the recognition result; The recognition results are converted into binary images and stored.

8. An image recognition model training system, characterized in that, include: The acquisition module is used to acquire the pre-trained dataset, the fine-tuned dataset, and the distillation dataset. The training module is used to train a pre-set pre-trained large model using the pre-trained dataset to obtain pre-trained weights. The fine-tuning module is used to initialize the preset fine-tuning model based on the pre-trained weights, and to fine-tune the initialized fine-tuning model using the fine-tuning dataset to obtain the teacher model; The distillation module is used to perform response distillation on a preset student model using the teacher model and the distillation dataset to obtain a target image recognition model.

9. An electronic device comprising a memory and a processor, wherein the memory stores a computer program, characterized in that, When the processor executes the computer program, it implements the image recognition model training method as described in any one of claims 1 to 7.

10. A computer-readable storage medium storing a computer program, characterized in that, When the computer program is executed by the processor, it implements the image recognition model training method as described in any one of claims 1 to 7.

Images