Image recognition model training method and apparatus, and electronic device
The image recognition model trained by pruning and loss value optimization solves the problem of high-precision real-time recognition and realizes fast and efficient image recognition on embedded devices.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- CHINA AUTOMOTIVE INNOVATION CORP
- Filing Date
- 2024-05-27
- Publication Date
- 2026-07-24
AI Technical Summary
Existing image recognition models are difficult to achieve high-precision real-time recognition on embedded devices with limited computing resources, and the model distillation process is time-consuming and inaccurate.
The student model is obtained by pruning the initial image recognition model. The loss value is determined by combining the real label data and the output feature map and prediction results of the teacher model. The student model is trained using the backpropagation algorithm to maintain structural similarity and achieve rapid convergence.
While maintaining high accuracy, the image recognition model is simplified, improving recognition speed and model training efficiency, making it suitable for devices with limited computing resources.
Smart Images
Figure CN118506127B_ABST
Abstract
Description
Technical Field
[0001] This application relates to the field of deep learning, and in particular to a training method, apparatus, and electronic device for an image recognition model. Background Technology
[0002] With the development of deep learning technology, image recognition models have been widely applied. For example, in the field of autonomous driving, image recognition models are often used for perception and decision-making to assist in achieving autonomous driving. However, high-precision image recognition models are often very large and complex in structure, requiring a lot of computation. In practical applications, complex image recognition models are difficult to meet the requirements of real-time tasks and cannot run on embedded devices with limited computing resources and memory space.
[0003] In practical applications, image recognition models are often distilled to compress them. Model distillation typically uses a pre-trained teacher model that performs well on large datasets, and then transfers the teacher model's knowledge to a student model, enabling the student model to perform inference on a smaller scale. However, existing model distillation schemes usually require numerous training iterations and computations, consuming significant time and computational resources, and the accuracy of the resulting student model often differs significantly from that of the teacher model.
[0004] Therefore, how to simplify image recognition models while maintaining accuracy is a technical problem that urgently needs to be solved. Summary of the Invention
[0005] This application provides a training method, apparatus, and electronic device for an image recognition model, which can simplify the image recognition model while maintaining accuracy.
[0006] On one hand, embodiments of this application provide a method for training an image recognition model, the method comprising:
[0007] Acquire training samples and an initial image recognition model, and designate the initial image recognition model as the teacher model; training samples include training images and training label data;
[0008] The teacher model is pruned according to a preset pruning rate to obtain the initial student model, and the initial student model is determined as the current student model;
[0009] Using the current student model, the training image is processed to obtain a first output feature map and a first prediction result; the first output feature map is the convolutional feature obtained after the training image is processed by the backbone network of the current student model; the first prediction result is the prediction result obtained after the first output feature map is processed by the prediction network of the current student model.
[0010] The loss value is determined based on the first output feature map, the first prediction result, the second output feature map, the second prediction result, and the training label data; the second output feature map is the convolutional feature obtained after the training image is processed by the backbone network of the teacher model; the second prediction result is the prediction result obtained after the second output feature map is processed by the prediction network of the teacher model.
[0011] The current student model is trained based on the loss value until the training termination condition is met, and the current student model at the end of training is determined as the target image recognition model.
[0012] In one exemplary embodiment, determining a loss value based on a first output feature map, a first prediction result, a second output feature map, a second prediction result, and training label data includes:
[0013] Based on the first prediction result and the training label data, determine the first loss;
[0014] Based on the first and second prediction results, the second loss is determined;
[0015] The third loss is determined based on the first and second output feature maps;
[0016] The loss value is determined based on the first loss, the second loss, and the third loss.
[0017] In one exemplary embodiment, determining the loss value based on the first loss, the second loss, and the third loss includes:
[0018] Determine the weights corresponding to the first loss, the second loss, and the third loss;
[0019] Based on the respective weights of the first loss, the second loss, and the third loss, the first loss, the second loss, and the third loss are weighted according to a preset strategy to obtain the loss value.
[0020] In one exemplary implementation, training the current student model based on the loss value until a training termination condition is met includes:
[0021] Based on the loss value, adjust the parameters of the current student model until the training termination condition is met. The training termination condition includes: the training reaches a preset number of times, or the loss value is less than a first threshold, or the loss value is equal to the first threshold.
[0022] In one exemplary embodiment, the method further includes:
[0023] Acquire test samples and teacher indicator data; test samples include test images and test label data, and teacher indicator data are the data corresponding to the model evaluation indicators of the teacher model on the test samples;
[0024] When the training termination condition is not met, the method also includes:
[0025] The current student model is evaluated using test samples to obtain student indicator data;
[0026] When the difference between student indicator data and teacher indicator data is greater than the second threshold, the parameters of the current student model are adjusted according to the loss value to obtain the adjusted student model, and the adjusted student model is re-determined as the current student model; return to the execution step: using the current student model, the training image is processed to obtain the first output feature map and the first prediction result;
[0027] When the difference between the student indicator data and the teacher indicator data is less than or equal to the second threshold, the training of the current student model ends and the current student model is determined as the target image recognition model.
[0028] In one exemplary embodiment, the current student model is evaluated using test samples to obtain student indicator data, including:
[0029] The test image is input into the current student model for processing to obtain the third prediction result;
[0030] The third prediction result is compared with the test label data in the test sample to obtain the student indicator data.
[0031] On the other hand, this application provides an image recognition method, which includes:
[0032] Acquire the image to be recognized;
[0033] The target image recognition model is used to process the image to be recognized to obtain the image recognition result. The target image recognition model is obtained according to the training method of the image recognition model described above.
[0034] On the other hand, this application provides a training apparatus for an image recognition model, the apparatus comprising:
[0035] The acquisition module is used to acquire training samples and an initial image recognition model, and to identify the initial image recognition model as the teacher model; the training samples include training images and training label data.
[0036] The pruning module is used to prune the teacher model according to a preset pruning rate to obtain the initial student model, and then determine the initial student model as the current student model.
[0037] The student output module is used to process the training image using the current student model to obtain a first output feature map and a first prediction result. The first output feature map is the convolutional feature obtained after the training image is processed by the backbone network of the current student model. The first prediction result is the prediction result obtained after the first output feature map is processed by the prediction network of the current student model.
[0038] The loss determination module is used to determine the loss value based on the first output feature map, the first prediction result, the second output feature map, the second prediction result, and the training label data; the second output feature map is the convolutional feature obtained after the training image is processed by the backbone network of the teacher model; the second prediction result is the prediction result obtained after the second output feature map is processed by the prediction network of the teacher model.
[0039] The training module is used to train the current student model based on the loss value until the training termination condition is met, and to determine the current student model at the end of training as the target image recognition model.
[0040] On the other hand, this application provides an image recognition device, which includes:
[0041] The acquisition module is used to acquire the image to be recognized;
[0042] The processing module is used to process the image to be recognized using the target image recognition model to obtain the image recognition result. The target image recognition model is obtained according to the training method of the image recognition model mentioned above.
[0043] On the other hand, this application also provides an electronic device, the device including a processor and a memory, the memory storing at least one instruction or at least one program, the at least one instruction or the at least one program being loaded and executed by the processor to implement the training method of the image recognition model as described above, or the image recognition method as described above.
[0044] On the other hand, this application also provides a computer storage medium storing at least one instruction or at least one program, which is loaded and executed by a processor to implement the training method of the image recognition model as described above, or the image recognition method as described above.
[0045] On the other hand, this application also provides a computer program product or computer program that includes 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 a training method for the image recognition model as described above, or an image recognition method as described above.
[0046] The training method for the image recognition model provided in this application has the following technical effects:
[0047] This application uses a pruned initial image recognition model as a student model, enabling rapid student model construction. Using the pruned model as the student model ensures structural similarity between the teacher and student models, facilitating the transfer of knowledge from the teacher model to the student model and ensuring the trained student model maintains the accuracy of the teacher model. Furthermore, this application's scheme determines the loss value by combining real-label data with the output feature maps and prediction results of both the student and teacher models, enabling rapid model convergence and accelerating the training process. Therefore, this method makes model training more efficient. Moreover, this training method simplifies the image recognition model while maintaining accuracy. Compared to the initial image recognition model, the image recognition model trained using this method has a more streamlined structure, effectively improving recognition speed while maintaining high accuracy. Attached Figure Description
[0048] To more clearly illustrate the technical solutions in the embodiments of this application, the drawings used in the description of the embodiments or the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of this application. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.
[0049] Figure 1 This is an application environment diagram of a training method for an image recognition model provided in an embodiment of this application.
[0050] Figure 2 This is a flowchart illustrating a training method for an image recognition model provided in an embodiment of this application.
[0051] Figure 3 This is a flowchart illustrating a method for determining a loss value provided in an embodiment of this application.
[0052] Figure 4 This is a flowchart illustrating a model training method provided in an embodiment of this application.
[0053] Figure 5 This is a schematic diagram of the training module structure of an image recognition model provided in an embodiment of this application.
[0054] Figure 6 This is a flowchart illustrating an image recognition method provided in an embodiment of this application.
[0055] Figure 7This is a schematic diagram of the structure of a training device for an image recognition model provided in an embodiment of this application.
[0056] Figure 8 This is a schematic diagram of the structure of an image recognition device provided in an embodiment of this application.
[0057] Figure 9 This is a hardware structure block diagram of a server for a training method of an image recognition model provided in an embodiment of this application. Detailed Implementation
[0058] The technical solutions of the embodiments of this application will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of this application, and not all embodiments. Based on the embodiments of this application, all other embodiments obtained by those of ordinary skill in the art without creative effort are within the scope of protection of this application.
[0059] It should be noted that the terms "first," "second," etc., in the specification, claims, and accompanying drawings of this application are used to distinguish similar objects and are not necessarily used to describe a specific order or sequence. It should be understood that such data can be interchanged where appropriate so that the embodiments of this application described herein can be implemented in orders other than those illustrated or described herein. Furthermore, the terms "comprising" and "having," and any variations thereof, are intended to cover non-exclusive inclusion; for example, a process, method, system, product, or server that comprises a series of steps or units is not necessarily limited to those steps or units explicitly listed, but may include other steps or units not explicitly listed or inherent to such processes, methods, products, or devices.
[0060] Figure 1 This is an application environment diagram of a training method for an image recognition model provided in an embodiment of this application.
[0061] Image recognition models can be applied to various scenarios depending on the needs. In the field of autonomous driving, image recognition models can be used to identify road scenes, such as lane lines, traffic lights, and vehicle types; in the field of home security, image recognition models can be used to identify pedestrians, items, and packages.
[0062] like Figure 1As shown, the initial image recognition model 11 is a pre-trained deep learning model that performs well on a large test dataset. The initial image recognition model 11 can be trained on server 10 using a large dataset. Server 10 can be a standalone server, a distributed server, or a server cluster consisting of multiple servers. Server 10 may include network communication units, processors, and memory, etc. In application, the initial image recognition model 11 can be migrated to terminal device 20, enabling terminal device 20 to perform image recognition functions. The target image recognition model 21 can be considered as the model obtained after migrating the initial image recognition model 11. Terminal device 20 may include physical devices such as smartphones, desktop computers, tablets, laptops, digital assistants, smart wearable devices, smart home devices, in-vehicle devices, and robots, and may also include software running on physical devices, such as web pages provided to users by service providers, or applications provided to users by these service providers.
[0063] However, since the initial image recognition model 11 is often complex in structure and has a large number of parameters, directly transferring it to the terminal device 20 would result in the terminal device 20 not having enough computing power to run it. At the same time, as the computational load of the model increases, the speed of the model in recognizing images will also slow down, and real-time performance cannot be guaranteed.
[0064] Therefore, this application proposes a training method for an image recognition model, which simplifies the image recognition model by distilling it while maintaining accuracy.
[0065] To better understand the solution of this application, the relevant terms and concepts in the technical solution of this application will be introduced below.
[0066] (1) Loss Function
[0067] During neural network training, to ensure the output closely approximates the desired predicted value, we compare the network's prediction with the target value. Based on the difference, we update the weight vectors of each layer (usually pre-configuring parameters before the initial update). For example, if the prediction is too high, the weight vectors are adjusted to predict a lower value. This adjustment continues until the model predicts the target value or a value very close to it. Therefore, we need to predefine "how to compare the difference between the predicted and target values," which is the loss function or objective function. These are crucial equations for measuring the difference between the predicted and target values. Taking the loss function as an example, a larger loss value indicates a greater difference, and training the neural network becomes a process of minimizing this loss.
[0068] (2) Backpropagation algorithm
[0069] Neural networks can employ backpropagation (BP) to correct the initial parameter values during training, thereby minimizing the reconstruction error loss. Specifically, forward propagation of the input signal to the output generates error loss; this error loss information is then propagated back to update the parameters of the initial neural network model, leading to convergence of the error loss. The backpropagation algorithm is an error-loss-driven backpropagation process aimed at obtaining the optimal parameters of the neural network model, such as the weight matrix.
[0070] Figure 2 This is a flowchart illustrating a training method for an image recognition model provided in an embodiment of this application.
[0071] S201: Obtain training samples and an initial image recognition model, and determine the initial image recognition model as the teacher model.
[0072] The training samples include training images and training label data.
[0073] The initial image recognition model is a large-scale image recognition model that has already been trained. The initial image recognition model can be an existing general-purpose image recognition model, or a model fine-tuned using specific training samples on top of a general-purpose model. This application does not specifically limit the type of initial image recognition model; any image recognition model that needs to be compressed can be used as an initial image recognition model. The initial image recognition model may include a backbone network and a prediction network. The backbone network is used to extract features, and the prediction network is used to output the prediction results.
[0074] It should be noted that using the initial image recognition model directly as the teacher model ensures that the teacher model has a certain guiding ability, which can better guide the student model learning.
[0075] Training samples include training images and training label data. Training label data can be information such as the target category and location contained in the training images, and can be labeled according to the actual application scenario.
[0076] S203: Prune the teacher model according to the preset pruning rate to obtain the initial student model, and determine the initial student model as the current student model.
[0077] Model pruning removes redundant connections and parameters from a model, reducing its parameters or structure. Compared to the teacher model, the student model is smaller. However, due to the removal of some parameters or structures, the model's feature representation ability is affected, and its performance decreases. The pruned model still needs to be trained. In this embodiment, the model obtained after pruning the teacher model is denoted as the initial student model, and then the initial student model is trained to improve the student model's recognition ability.
[0078] In one implementation, the teacher model is pruned according to a preset pruning rate, which can be done through methods such as parametric pruning, channel pruning, and structural pruning.
[0079] In one embodiment, pruning the teacher model according to a preset pruning rate to obtain an initial student model may include: determining the pruning position based on the structure of the teacher model; and pruning at the pruning position according to the preset pruning rate to obtain the initial student model. The pruning position may be some neural network layers of the model, some channels of the model, or a single neuron, etc., which can be determined according to the actual pruning algorithm.
[0080] For example, the DepGraph general structured pruning method can be used, with a preset pruning rate of 40%. Except for the output of the last layer of the teacher model (i.e., the prediction network), which remains unpruned, all other layers are pruned according to the preset pruning rate. It should be understood that the preset pruning rate value here is merely an example, and the preset pruning rate can be set according to actual needs.
[0081] S205: Using the current student model, process the training images to obtain the first output feature map and the first prediction result;
[0082] Wherein, the first output feature map is the convolutional feature obtained after the training image is processed by the backbone network of the current student model; the first prediction result is the prediction result obtained after the first output feature map is processed by the prediction network of the current student model.
[0083] By inputting the training image into the current student model, the prediction result of the current student model can be obtained. Specifically, the training image first undergoes feature extraction through the backbone network of the current student model, which outputs a first feature map. Then, the prediction network of the current student model processes the first feature map to obtain a first prediction result; this process can be viewed as a classification process. The first prediction result can include information such as the category, location, and confidence level of the identified target.
[0084] It should be noted that before inputting the training images into the current student model, the training images can be preprocessed, including image cropping and image pixel augmentation.
[0085] S207: Determine the loss value based on the first output feature map, the first prediction result, the second output feature map, the second prediction result, and the training label data.
[0086] The second output feature map is the convolutional feature obtained after the training image is processed by the backbone network of the teacher model; the second prediction result is the prediction result obtained after the second output feature map is processed by the prediction network of the teacher model.
[0087] Before this step, similar to step S205, the teacher model also processes the training images, and the output feature map is denoted as the second output feature map, and the output prediction result is denoted as the second prediction result. It should be noted that the training samples of the current student model and teacher model are the same.
[0088] In one implementation, the teacher model processes the training images after step S201, that is, immediately after obtaining the training samples, and then stores the second output feature map and the second prediction result. Then, it is read from the storage space when the loss value is determined in step S207.
[0089] In another implementation, the teacher model's processing of the training images can be performed simultaneously with step S205. That is, the training images are simultaneously input into the current student and teacher models for processing, resulting in a first output feature map, a first prediction result, a second output feature map, and a second prediction result.
[0090] Then, the loss value is determined based on the first output feature map, the first prediction result, the second output feature map, the second prediction result, and the training label data in the training samples. Compared with the training scheme in the prior art, the embodiments of this application introduce an intermediate layer output of the first and second output feature maps when determining the loss value, which can accelerate the convergence speed of the model and speed up the training process.
[0091] In one implementation, a loss function can be pre-constructed, and then the first output feature map, the first prediction result, the second output feature map, the second prediction result, and the training label data can be input into the loss function to obtain the loss value.
[0092] S209: Train the current student model based on the loss value until the training termination condition is met, and determine the current student model at the end of training as the target image recognition model.
[0093] Based on the loss value, the current student model can be trained using the backpropagation algorithm. The model obtained at the end of training is the target image recognition model. Compared to the initial image recognition model, this model is smaller, faster, and maintains essentially the same accuracy.
[0094] This application uses a pruned initial image recognition model as a student model, enabling rapid student model construction. Using the pruned model as the student model ensures structural similarity between the teacher and student models, facilitating the transfer of knowledge from the teacher model to the student model and ensuring the trained student model maintains the accuracy of the teacher model. Furthermore, this application's scheme determines the loss value by combining real-label data with the output feature maps and prediction results of both the student and teacher models, enabling rapid model convergence and accelerating the training process. Therefore, this method makes model training more efficient. Moreover, this training method simplifies the image recognition model while maintaining accuracy. Compared to the initial image recognition model, the image recognition model trained using this method has a more streamlined structure, effectively improving recognition speed while maintaining high accuracy.
[0095] Figure 3 This is a flowchart illustrating a method for determining a loss value provided in an embodiment of this application. Figure 3 This can be seen as a specific example of the method shown in step S207.
[0096] like Figure 3 As shown, determining the loss value based on the first output feature map, the first prediction result, the second output feature map, the second prediction result, and the training label data may include the following steps:
[0097] S301: Determine the first loss based on the first prediction result and the training label data;
[0098] The first loss represents the cross-entropy loss between the predicted output of the current student model and the true training labels. In one implementation, the first loss can be determined by the following formula:
[0099]
[0100] Where CE(y,p) represents the first loss, y represents the training label data, p represents the first prediction result, and n represents the number of classification categories.
[0101] S303: Determine the second loss based on the first and second prediction results;
[0102] The second loss represents the cross-entropy loss between the predictions output by the current student model and the predictions output by the teacher model. In one implementation, the second loss can be determined by the following formula:
[0103]
[0104] Where CE(q,p) represents the second loss, p represents the first prediction result, q represents the second prediction result, i.e. the prediction result output by the teacher model, and n represents the number of classification categories.
[0105] S305: Determine the third loss based on the first output feature map and the second output feature map;
[0106] The third loss represents the similarity loss between the output feature maps of the current student model's backbone network and the output feature maps of the teacher model's backbone network. In one implementation, the third loss can be determined by the following formula:
[0107]
[0108] Where Similarity(y1,y2) represents the third loss, namely the cosine structure similarity between the first output feature map y1 and the second output feature map y2, A i B represents the i-th feature value in the first output feature map y1. i Let y represent the i-th feature value in the second output feature map y2, and n represent the number of feature values in the first output feature map y1 and the second output feature map y2.
[0109] S307: Determine the loss value based on the first loss, the second loss, and the third loss.
[0110] In one implementation, determining the loss value based on the first loss, the second loss, and the third loss includes: determining the weights corresponding to the first loss, the second loss, and the third loss respectively; and weighting the first loss, the second loss, and the third loss according to a preset strategy based on the weights corresponding to the first loss, the second loss, and the third loss respectively to obtain the loss value.
[0111] In this embodiment, the loss function formula is as follows:
[0112] Loss total =α*CE(y,p)+β*CE(q,p)+γ*(1-Similarity(y1,y2))
[0113] Among them, Loss total Let represent the total loss, CE(y,p) represent the first loss, CE(q,p) represent the second loss, Similarity(y1,y2) represent the third loss, and α, β, and γ represent the weights corresponding to the first, second, and third losses, respectively.
[0114] The weights corresponding to the first loss, the second loss, and the third loss can be set according to actual needs. Preferably, the weights α, β, and γ corresponding to the first loss, the second loss, and the third loss can be 0.4, 0.4, and 0.2, respectively.
[0115] In another implementation, the loss value can be determined by directly adding the first loss, the second loss, and the third loss.
[0116] In one embodiment, training the current student model based on the loss value until a training termination condition is met may include: adjusting the parameters of the current student model based on the loss value until the training termination condition is met. The training termination condition includes: reaching a preset number of training iterations, or the loss value being less than a first threshold, or the loss value being equal to the first threshold. In one example, when the training termination condition is reaching a preset number of training iterations, assuming 100 epochs are trained, the student model obtained after 100 iterations is determined as the target image recognition model.
[0117] In one embodiment, test samples and teacher indicator data are first acquired. The test samples include test images and test label data, while the teacher indicator data refers to the model evaluation metrics corresponding to the teacher model on the test samples. After each training iteration, the test samples can be used to obtain the student indicator data for the current student model. It should be noted that the test samples contain multiple test images and corresponding test label data.
[0118] When the training termination criteria are not met, the training method may also include:
[0119] The current student model is evaluated using test samples to obtain student indicator data;
[0120] Specifically, this step may include: inputting the test image into the current student model for processing to obtain a third prediction result; comparing the third prediction result with the test label data in the test samples to obtain student indicator data. It should be understood that the indicator data can be commonly used model evaluation metrics, such as accuracy, recall, etc.
[0121] When the difference between student indicator data and teacher indicator data exceeds a second threshold, the parameters of the current student model are adjusted based on the loss value to obtain an adjusted student model, which is then redefined as the current student model. The process returns to the previous step: using the current student model, the training images are processed to obtain the first output feature map and the first prediction result. When the difference between student indicator data and teacher indicator data exceeds the second threshold, it indicates that the performance of the current student model differs significantly from that of the teacher model, and further training is required. The second threshold can be a difference threshold for a single indicator or a combined difference threshold for multiple indicators. The specific value of the second threshold can be set according to actual application requirements; this application does not limit the specific value of the second threshold.
[0122] When the difference between the student's metric data and the teacher's metric data is less than or equal to the second threshold, training of the current student model ends, and the current student model is designated as the target image recognition model. When the difference between the student's metric data and the teacher's metric data is less than or equal to the second threshold, it indicates that the performance of the current student model is close to that of the teacher model, and training can end.
[0123] When the student's performance metrics exceed the teacher's performance metrics, training of the current student model can be terminated, and the current student model can be designated as the target image recognition model. When the student's performance metrics exceed the teacher's performance metrics, it indicates that the current student model outperforms the teacher model, and training can be terminated.
[0124] In one embodiment, accuracy can be used as the evaluation metric. If the accuracy of the current student model is less than that of the teacher model, and the difference between the two is greater than a second threshold, the parameters of the current student model are adjusted based on the loss value. If the difference in accuracy between the current student model and the teacher model is less than or equal to the second threshold, training of the current student model ends, and the current student model is identified as the target image recognition model. If the accuracy of the current student model is greater than that of the teacher model, training of the current student model ends, and the current student model is identified as the target image recognition model. For example, the second threshold can be set to 1. Assuming the accuracy of the teacher model is 90, if after one iteration of training the current student model's accuracy is 86, the difference is 4, which is greater than the second threshold, then training continues. If after one iteration of training the current student model's accuracy is 89.4, the difference is 0.6, which is less than the second threshold, then training can end. If after one iteration of training the current student model's accuracy is 90.2, which is greater than the teacher model's accuracy, then training can end.
[0125] Figure 4 This is a flowchart illustrating a model training method provided in an embodiment of this application. Figure 4 This can be seen as a specific example of the method shown in step S209.
[0126] like Figure 4 As shown in (A) of this application, a model training method provided in this embodiment includes the following steps:
[0127] S401: Get data.
[0128] The data obtained in step S401 can be seen as a specific example of the training samples in step S201 above.
[0129] S402: Data preprocessing.
[0130] The acquired data is preprocessed, including image cropping and augmentation, to meet the model's input size requirements.
[0131] S403: Construct the pruning distillation training module.
[0132] This step can be seen as a specific example of steps S203 and S205.
[0133] S404: Determine whether training is complete.
[0134] The completion of training can be determined based on whether the training termination condition is met. Training ends when the termination condition is met, and the current student model is designated as the target image recognition model. If the termination condition is not met, step S405 is executed.
[0135] S405: Test set evaluation.
[0136] Among them, such as Figure 4 As shown in (B) above, this step may include the following sub-steps:
[0137] S4051: Get data.
[0138] The data obtained in step S4051 can be seen as a specific example of the test sample in the above embodiments.
[0139] S4052: Data preprocessing.
[0140] The acquired data is preprocessed, including image cropping and augmentation, to meet the model's input size requirements.
[0141] S4053: Pruned Distillation Student Model Reasoning.
[0142] The pruned distillation student model in this step can be seen as a specific example of the current student model in the above embodiments. The reasoning process of the pruned distillation student model is described in the above embodiments and will not be repeated here.
[0143] S4055: Model post-processing.
[0144] Model post-processing refers to integrating the prediction results output by the pruned distillation student model and calculating indicator data, such as accuracy and recall.
[0145] S406: Determine whether the indicator is met.
[0146] The metric in step S406 can be seen as a specific example of the student metric data in the above embodiments. When the metric is close to or exceeds the metric data of the teacher model, training ends, and the current student model is determined as the target image recognition model. When the metric differs significantly from the metric data of the teacher model, training continues, and step S404 is executed.
[0147] Figure 5 This is a schematic diagram of the training module structure of an image recognition model provided in an embodiment of this application.
[0148] like Figure 5As shown, the teacher model consists of a backbone network and a prediction network. The output of the last layer of the backbone network is denoted as y2, and the prediction result of the prediction network is denoted as q. The student model also consists of a backbone network and a prediction network. The student model's backbone network is obtained by performing DepGraph structured pruning on the teacher model. The output of the last layer of the student model's backbone network is denoted as y1, and the prediction result of the prediction network is denoted as p. The student network is then trained. The loss function during training includes CE(y,p), which is the cross-entropy loss between the predicted result of the student model and the true training label; CE(q,p), which is the cross-entropy loss between the predicted result of the student model and the predicted result of the teacher model; and Similarity(y1,y2), which is the cosine structural similarity loss between the feature maps output by the student model's backbone network and the feature maps output by the teacher model's backbone network. This loss can be used as a relay supervision loss, which can accelerate model convergence.
[0149] Table 1 shows the test results of the target image recognition model trained using the training method provided in the above embodiments.
[0150] Table 1
[0151]
[0152] Taking ResNet-50 as the initial image recognition model and 100 training iterations as an example, the test results obtained from the ablation experiment are shown in Table 1. The accuracy metric can be equated to the model's evaluation metric acc1 on the ImageNet-1 test data, and the optimal model epoch is the number of epochs required to train the optimal model. The target model is the target image recognition model trained using the image recognition model training method provided in this application.
[0153] As shown in Table 1, the teacher model has the largest number of parameters and multiply-accumulate operations, indicating that it is the most complex model and takes the longest time to process a single image. Correspondingly, it also has the highest accuracy. Among the initial student model, the model trained with a loss function without a third loss, and the target model, the target model has the highest accuracy, is closest to the teacher model, and requires the fewest epochs to obtain the optimal model, indicating the fastest convergence speed during training. Furthermore, compared with the model trained with a loss function without a third loss, it can be shown that the third loss used in determining the loss function in this application is beneficial to the rapid convergence of the model.
[0154] Figure 6 This is a flowchart illustrating an image recognition method provided in an embodiment of this application.
[0155] like Figure 6As shown in the embodiment of this application, an image recognition method includes:
[0156] S601: Acquire the image to be recognized.
[0157] S603: Process the image to be recognized using the target image recognition model to obtain the image recognition result.
[0158] The target image recognition model is obtained according to any of the image recognition model training methods provided in the above embodiments. In the field of autonomous driving, the target image recognition model can improve the real-time performance of autonomous driving systems, reduce the demand for computing resources, and help realize low-power autonomous driving solutions on embedded devices.
[0159] As can be seen from the technical solutions provided in the above embodiments of this application, this application uses the pruned initial image recognition model as the student model, thus achieving rapid construction of the student model. Using the pruned model as the student model ensures the structural similarity between the teacher model and the student model, which is more conducive to transferring the knowledge of the teacher model to the student model, enabling the trained student model to maintain the accuracy of the teacher model. Furthermore, the solution of this application determines the loss value by combining the real label data with the output feature maps and prediction results of the student model and the teacher model. In particular, the introduction of a third loss can achieve rapid model convergence and accelerate the model training process. Therefore, the method of this application can make the model training process more efficient. Moreover, the training method of this application can simplify the image recognition model while maintaining accuracy. Compared with the initial image recognition model, the image recognition model trained by the method of this application has a more streamlined structure, which can effectively improve the recognition speed while maintaining high accuracy.
[0160] Figure 7 This is a schematic diagram of the structure of a training device for an image recognition model provided in an embodiment of this application.
[0161] like Figure 7 As shown, the device 1000 includes:
[0162] The acquisition module 1001 is used to acquire training samples and an initial image recognition model, and to determine the initial image recognition model as the teacher model; the training samples include training images and training label data.
[0163] The pruning module 1002 is used to prune the teacher model according to a preset pruning rate to obtain the initial student model, and to determine the initial student model as the current student model.
[0164] The student output module 1003 is used to process the training image using the current student model to obtain a first output feature map and a first prediction result; the first output feature map is the convolutional feature obtained after the training image is processed by the backbone network of the current student model; the first prediction result is the prediction result obtained after the first output feature map is processed by the prediction network of the current student model.
[0165] The loss determination module 1004 is used to determine the loss value based on the first output feature map, the first prediction result, the second output feature map, the second prediction result, and the training label data; the second output feature map is the convolutional feature obtained after the training image is processed by the backbone network of the teacher model; the second prediction result is the prediction result obtained after the second output feature map is processed by the prediction network of the teacher model.
[0166] The training module 1005 is used to train the current student model based on the loss value until the training termination condition is met, and to determine the current student model at the end of training as the target image recognition model.
[0167] In some embodiments, the loss determination module may include:
[0168] The first loss determination submodule is used to determine the first loss based on the first prediction result and the training label data;
[0169] The second loss determination submodule is used to determine the second loss based on the first prediction result and the second prediction result;
[0170] The third loss determination submodule is used to determine the third loss based on the first output feature map and the second output feature map;
[0171] The loss determination submodule is used to determine the loss value based on the first loss, the second loss, and the third loss.
[0172] In some embodiments, the loss determination submodule may include:
[0173] The weight determination unit is used to determine the weights corresponding to the first loss, the second loss, and the third loss, respectively.
[0174] The weighting unit is used to weight the first loss, the second loss, and the third loss according to their respective weights and a preset strategy to obtain the loss value.
[0175] In some embodiments, the training module may include:
[0176] The parameter adjustment submodule is used to adjust the parameters of the current student model according to the loss value until the training termination condition is met. The training termination condition includes: the training reaches a preset number of times, or the loss value is less than a first threshold, or the loss value is equal to the first threshold.
[0177] In some embodiments, the device may further include:
[0178] The test data acquisition submodule is used to acquire test samples and teacher indicator data; the test samples include test images and test label data, and the teacher indicator data is the data corresponding to the model evaluation indicators of the teacher model on the test samples;
[0179] When the training termination condition is not met, the device may further include:
[0180] The evaluation unit is used to evaluate the current student model using test samples and obtain student indicator data;
[0181] The adjustment unit is used to adjust the parameters of the current student model according to the loss value when the difference between the student indicator data and the teacher indicator data is greater than the second threshold, so as to obtain the adjusted student model and redetermine the adjusted student model as the current student model; return to the execution step: using the current student model, process the training image to obtain the first output feature map and the first prediction result;
[0182] The model determination unit is used to terminate the training of the current student model and determine the current student model as the target image recognition model when the difference between the student indicator data and the teacher indicator data is less than or equal to the second threshold.
[0183] In some embodiments, the evaluation unit may include:
[0184] The test image processing unit is used to input the test image into the current student model for processing and to obtain a third prediction result.
[0185] The student indicator determination unit is used to compare the third prediction result with the test label data in the test sample to obtain student indicator data.
[0186] Figure 8 This is a schematic diagram of the structure of an image recognition device provided in an embodiment of this application.
[0187] like Figure 8 As shown, the device 2000 includes:
[0188] The acquisition module 2001 is used to acquire the image to be recognized;
[0189] The processing module 2002 is used to process the image to be recognized using a target image recognition model to obtain the image recognition result. The target image recognition model is obtained according to the training method of the image recognition model provided in the embodiments of this application.
[0190] The apparatus and method embodiments described herein are based on the same inventive concept.
[0191] This application provides an electronic device including a processor and a memory. The memory stores at least one instruction or at least one program, which is loaded and executed by the processor to implement the method provided in the above method embodiments.
[0192] Embodiments of this application also provide a computer storage medium, which can be disposed in a terminal to store at least one instruction or at least one program related to implementing a method as provided in the above method embodiments, wherein the at least one instruction or at least one program is loaded and executed by the processor to implement the method provided in the above method embodiments.
[0193] Embodiments of this application also provide a computer program product or computer program, which includes 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 a training method for the image recognition model provided in the above-described method embodiments, or an image recognition method as provided in the above-described method embodiments.
[0194] Optionally, in this embodiment, the storage medium may be located at at least one of the multiple network servers in a computer network. Optionally, in this embodiment, the storage medium may include, but is not limited to, various media capable of storing program code, such as USB flash drives, read-only memory (ROM), random access memory (RAM), portable hard drives, magnetic disks, or optical disks.
[0195] The memory described in this application embodiment can be used to store software programs and modules. The processor executes various functional applications and data processing by running the software programs and modules stored in the memory. The memory may mainly include a program storage area and a data storage area. The program storage area may store the operating system, applications required for the functions, etc.; the data storage area may store data created according to the use of the device, etc. In addition, the memory may include high-speed random access memory, and may also include non-volatile memory, such as at least one disk storage device, flash memory device, or other volatile solid-state storage device. Accordingly, the memory may also include a memory controller to provide the processor with access to the memory.
[0196] The image recognition model training method provided in this application can be executed on a mobile terminal, computer terminal, server, or similar computing device. Taking running on a server as an example, Figure 9This is a hardware structure block diagram of a server for a training method of an image recognition model provided in an embodiment of this application. For example... Figure 9 As shown, server 900 can be considered a specific example of an electronic device. Server 900 can vary considerably due to different configurations or performance, and may include one or more central processing units (CPUs) 910 (CPUs 910 may include, but are not limited to, microprocessors (MCUs) or programmable logic devices (FPGAs), a memory 930 for storing data, and one or more storage media 920 (e.g., one or more mass storage devices) for storing application programs 923 or data 922. The memory 930 and storage media 920 can be temporary or persistent storage. The program stored in storage media 920 may include one or more modules, each module including a series of instruction operations on the server. Furthermore, the CPU 910 may be configured to communicate with storage media 920 and execute the series of instruction operations in storage media 920 on server 900. Server 900 may also include one or more power supplies 960, one or more wired or wireless network interfaces 950, one or more input / output interfaces 940, and / or one or more operating systems 921, such as Windows Server™, Mac OS X™, Unix™, Linux™, FreeBSD™, etc.
[0197] The input / output interface 940 can be used to receive or send data via a network. Specific examples of the network described above may include a wireless network provided by the communication provider of server 900. In one example, the input / output interface 940 includes a Network Interface Controller (NIC), which can connect to other network devices via a base station to communicate with the Internet. In another example, the input / output interface 940 may be a radio frequency (RF) module for wireless communication with the Internet.
[0198] Those skilled in the art will understand that Figure 9 The structure shown is for illustrative purposes only and does not limit the structure of the aforementioned electronic device. For example, server 900 may also include... Figure 9 The more or fewer components shown, or having the same Figure 9 The different configurations shown.
[0199] As can be seen from the embodiments of the image recognition model training method, apparatus, and electronic device provided in this application, this application uses the pruned initial image recognition model as the student model, achieving rapid construction of the student model, ensuring structural similarity between the teacher model and the student model, and facilitating the transfer of knowledge from the teacher model to the student model. The scheme of this application determines the loss value by combining real label data with the output feature maps and prediction results of the student and teacher models, enabling rapid model convergence and accelerating the model training process. Therefore, the method of this application makes the model training process more efficient. Furthermore, the training method of this application simplifies the image recognition model while maintaining accuracy. Compared to the initial image recognition model, the image recognition model trained by the method of this application has a more streamlined structure, effectively improving recognition speed while maintaining high accuracy.
[0200] It should be noted that the order of the embodiments described above is merely for descriptive purposes and does not represent the superiority or inferiority of the embodiments. Furthermore, specific embodiments have been described above. Other embodiments are within the scope of the appended claims. In some cases, the actions or steps described in the claims can be performed in a different order than that shown in the embodiments and still achieve the desired result. Additionally, the processes depicted in the drawings do not necessarily require a specific or sequential order to achieve the desired result. In some embodiments, multitasking and parallel processing are also possible or may be advantageous.
[0201] The various embodiments in this specification are described in a progressive manner. Similar or identical parts between embodiments can be referred to mutually. Each embodiment focuses on describing the differences from other embodiments. In particular, the embodiments of apparatus, devices, and storage media are basically similar to the method embodiments, so the descriptions are relatively simple; relevant parts can be referred to the descriptions of the method embodiments.
[0202] Those skilled in the art will understand that all or part of the steps of the above embodiments can be implemented by hardware or by a program instructing related hardware. The program can be stored in a computer storage medium, such as a read-only memory, a disk, or an optical disk.
[0203] The above description is only a preferred embodiment of this application and is not intended to limit this application. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of this application should be included within the protection scope of this application.
Claims
1. A training method for an image recognition model, characterized in that, include: Acquire training samples and an initial image recognition model, and determine the initial image recognition model as the teacher model; the training samples include training images and training label data; The teacher model is pruned according to a preset pruning rate to obtain an initial student model, and the initial student model is determined as the current student model; The training image is processed using the current student model to obtain a first output feature map and a first prediction result; the first output feature map is the convolutional feature obtained after the training image is processed by the backbone network of the current student model. The first prediction result is the prediction result obtained after the first output feature map is processed by the prediction network of the current student model; The loss value is determined based on the first output feature map, the first prediction result, the second output feature map, the second prediction result, and the training label data; the second output feature map is the convolutional feature obtained after the training image is processed by the backbone network of the teacher model; the second prediction result is the prediction result obtained after the second output feature map is processed by the prediction network of the teacher model. The current student model is trained according to the loss value until the training termination condition is met, and the current student model at the end of training is determined as the target image recognition model. The training termination condition includes: the training reaches a preset number of times, or the loss value is less than a first threshold, or the loss value is equal to the first threshold.
2. The method according to claim 1, characterized in that, The step of determining the loss value based on the first output feature map, the first prediction result, the second output feature map, the second prediction result, and the training label data includes: Based on the first prediction result and the training label data, a first loss is determined; Based on the first prediction result and the second prediction result, determine the second loss; Based on the first output feature map and the second output feature map, a third loss is determined; The loss value is determined based on the first loss, the second loss, and the third loss.
3. The method according to claim 1, characterized in that, The step of training the current student model based on the loss value until the training termination condition is met includes: Based on the loss value, adjust the parameters of the current student model until the training termination condition is met.
4. The method according to claim 1, characterized in that, The method further includes: Acquire test samples and teacher indicator data; the test samples include test images and test label data, and the teacher indicator data are the data corresponding to the model evaluation indicators of the teacher model on the test samples; When the training termination condition is not met, the method further includes: The current student model is evaluated using the test samples to obtain student indicator data; When the difference between the student indicator data and the teacher indicator data is greater than the second threshold, the parameters of the current student model are adjusted according to the loss value to obtain the adjusted student model, and the adjusted student model is re-determined as the current student model; return to the execution step: use the current student model to process the training image to obtain the first output feature map and the first prediction result; When the difference between the student indicator data and the teacher indicator data is less than or equal to the second threshold, the training of the current student model ends and the current student model is determined as the target image recognition model.
5. The method according to claim 4, characterized in that, The step of evaluating the current student model using the test samples to obtain student indicator data includes: The test image is input into the current student model for processing to obtain a third prediction result; The third prediction result is compared with the test label data in the test sample to obtain student indicator data.
6. An image recognition method, characterized in that, include: Acquire the image to be recognized; The image to be recognized is processed using a target image recognition model to obtain an image recognition result. The target image recognition model is obtained by the training method of the image recognition model according to any one of claims 1 to 5.
7. A training device for an image recognition model, characterized in that, include: The acquisition module is used to acquire training samples and an initial image recognition model, and to determine the initial image recognition model as the teacher model; The training samples include training images and training label data; The pruning module is used to prune the teacher model according to a preset pruning rate to obtain an initial student model, and to determine the initial student model as the current student model; The student output module is used to process the training image using the current student model to obtain a first output feature map and a first prediction result; the first output feature map is the convolutional feature obtained after the training image is processed by the backbone network of the current student model; The first prediction result is the prediction result obtained after the first output feature map is processed by the prediction network of the current student model; The loss determination module is used to determine a loss value based on the first output feature map, the first prediction result, the second output feature map, the second prediction result, and the training label data; the second output feature map is the convolutional feature obtained after the training image is processed by the backbone network of the teacher model; The second prediction result is the prediction result obtained after the second output feature map is processed by the prediction network of the teacher model; The training module is used to train the current student model according to the loss value until the training termination condition is met, and to determine the current student model at the end of training as the target image recognition model. The training termination condition includes: the training reaches a preset number of times, or the loss value is less than a first threshold, or the loss value is equal to the first threshold.
8. An image recognition device, characterized in that, include: The acquisition module is used to acquire the image to be recognized; The processing module is used to process the image to be recognized using a target image recognition model to obtain an image recognition result, wherein the target image recognition model is obtained by the training method of the image recognition model according to any one of claims 1 to 5.
9. An electronic device, characterized in that, The electronic device includes a processor and a memory, wherein the memory stores at least one instruction or at least one program, the at least one instruction or at least one program being loaded and executed by the processor to implement the training method of the image recognition model as described in any one of claims 1-5, or the image recognition method as described in claim 6.
10. A computer storage medium, characterized in that, The computer storage medium stores at least one instruction or at least one program, which is loaded and executed by a processor to implement the training method of the image recognition model as described in any one of claims 1-5, or the image recognition method as described in claim 6.