Multi-task annotation model training method and device, electronic equipment and storage medium

Abstract

The invention relates to the technical field of artificial intelligence, and provides a multi-task annotation model training method and device, electronic equipment and a storage medium. According to the method, a training sample and a pre-training multi-task annotation model are obtained, a shared feature extraction network in the pre-training multi-task annotation model is firstly trained by using the training sample, and network parameters of the shared feature extraction network are frozen after training of the shared feature extraction network is completed; and then training the pre-trained multi-task labeling model in which the network parameters of the shared feature extraction network are frozen by using the training sample again to obtain N trained classification networks and further obtain the trained multi-task labeling model, so that the labeling model capable of identifying various different types of data can be obtained by training only one model, and the accuracy of labeling is improved. Meanwhile, parameter freezing can be carried out on the shared feature extraction network after training is completed, network parameters of the shared feature extraction network do not need to be adjusted when all classification networks are trained subsequently, the data processing amount is reduced, the consistency of labeling results can be well kept, and the data labeling efficiency is improved.

Description

Technical Field

[0001] The present application relates to the field of artificial intelligence technology, and in particular to a training method, device, electronic device and storage medium for a multi-task labeling model. Background Art

[0002] In the process of pre-labeling of autonomous driving data, labeling of long-tail data is a common and difficult challenge. In the related technology, when the labeling effect of long-tail data is not good, either update the samples including the data of the relevant type of long-tail data, retrain the labeling model, or try to enhance the long-tail data, or adjust the category weight of the long-tail data to solve the problem.

[0003] These processing methods are costly, and in some cases the trained annotation models are still not accurate enough for annotating long-tail data. Summary of the invention

[0004] In view of this, an embodiment of the present application provides a training method, device, electronic device and storage medium for a multi-task labeling model to solve the problem of inaccurate labeling of long-tail data in the prior art.

[0005] A first aspect of an embodiment of the present application provides a training method for a multi-task annotation model, comprising:

[0006] Acquire a training sample, wherein the training sample includes an image data set, each image data in the image data set is respectively marked with at least one category label, and the image data set corresponds to N different categories, where N is a positive integer greater than 1;

[0007] Obtain a pre-trained multi-task annotation model, wherein the pre-trained multi-task annotation model includes a shared feature extraction network and N different classification networks, each classification network corresponding to a category;

[0008] training the pre-trained multi-task annotation model using the training samples, and freezing network parameters of the shared feature extraction network in response to determining that the shared feature extraction network has converged;

[0009] The pre-trained multi-task labeling model with the network parameters of the shared feature extraction network frozen is trained again using the training samples until all the N classification networks converge, thereby obtaining a trained multi-task labeling model.

[0010] A second aspect of the embodiments of the present application provides a data labeling method, including:

[0011] Get target data;

[0012] Obtain a multi-task annotation model trained by the method provided in the first aspect above;

[0013] Using the trained multi-task annotation model to predict the target data, to obtain a predicted classification;

[0014] The target data is labeled using the predicted classification.

[0015] According to a third aspect of an embodiment of the present application, a training device for a multi-task annotation model is provided, comprising:

[0016] An acquisition module is configured to acquire a training sample, wherein the training sample includes an image data set, each image data in the image data set is respectively marked with at least one category label, and the image data set corresponds to N different categories, where N is a positive integer greater than 1;

[0017] The acquisition module is further configured to acquire a pre-trained multi-task annotation model, wherein the pre-trained multi-task annotation model includes a shared feature extraction network and N different classification networks, each classification network corresponding to a category;

[0018] a training module configured to train the pre-trained multi-task annotation model using the training samples, and in response to determining that the shared feature extraction network has converged, freeze network parameters of the shared feature extraction network;

[0019] The training module is also configured to use the training samples to retrain the pre-trained multi-task labeling model with the network parameters of the shared feature extraction network frozen until the N classification networks converge to obtain a trained multi-task labeling model.

[0020] According to a fourth aspect of an embodiment of the present application, an electronic device is provided, comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor implements the steps of the above method when executing the computer program.

[0021] According to a fifth aspect of an embodiment of the present application, a computer-readable storage medium is provided, which stores a computer program, and when the computer program is executed by a processor, the steps of the above method are implemented.

[0022] Compared with the prior art, the embodiments of the present application have the following beneficial effects: the embodiments of the present application obtain training samples and a pre-trained multi-task labeling model, use the training samples to first train the shared feature extraction network in the pre-trained multi-task labeling model, and freeze its network parameters after the shared feature extraction network training is completed, and then use the training samples to train the pre-trained multi-task labeling model with the network parameters of the shared feature extraction network frozen again, so as to obtain N trained classification networks, and then obtain the trained multi-task labeling model, so as to achieve the goal of obtaining a labeling model that can recognize multiple different categories of data by training only one model, and at the same time, the shared feature extraction network can be parameter-frozen after training, and there is no need to adjust the network parameters of the shared feature extraction network when subsequently training each classification network, thereby reducing the amount of data processing and being able to better maintain the consistency of the labeling results, thereby improving the efficiency of data labeling. BRIEF DESCRIPTION OF THE DRAWINGS

[0023] In order to more clearly illustrate the technical solutions in the embodiments of the present application, the drawings required for use in the embodiments or the description of the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present application. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative work.

[0024] Figure 1 It is a flowchart of a training method for a multi-task labeling model provided in an embodiment of the present application.

[0025] Figure 2 It is a flowchart of a method for determining convergence of a shared feature extraction network provided in an embodiment of the present application.

[0026] Figure 3 It is a flowchart of a method provided in an embodiment of the present application for retraining a pre-trained multi-task labeling model that has frozen network parameters of a shared feature extraction network using training samples until all N classification networks converge.

[0027] Figure 4 It is a flowchart of a method for retraining a multi-task labeling model using updated training samples provided in an embodiment of the present application.

[0028] Figure 5 It is a flowchart of a data labeling method provided in an embodiment of the present application.

[0029] Figure 6 It is a schematic diagram of a training device for a multi-task labeling model provided in an embodiment of the present application.

[0030] Figure 7 It is a schematic diagram of an electronic device provided in an embodiment of the present application. DETAILED DESCRIPTION

[0031] In the following description, specific details such as specific system structures, technologies, etc. are provided for the purpose of illustration rather than limitation, so as to provide a thorough understanding of the embodiments of the present application. However, it should be clear to those skilled in the art that the present application may also be implemented in other embodiments without these specific details. In other cases, detailed descriptions of well-known systems, devices, circuits, and methods are omitted to prevent unnecessary details from obstructing the description of the present application.

[0032] A data annotation method and device according to an embodiment of the present application will be described in detail below with reference to the accompanying drawings.

[0033] As mentioned above, when the relevant technology encounters a situation where the long-tail data labeling effect is not good, either the samples including the long-tail data-related type data are updated, the labeling model is retrained, or the long-tail data is enhanced, or the category weight of the long-tail data is adjusted to solve the problem.

[0034] That is to say, in related technologies, a unified model is usually used to train all categories, and then the model is used for labeling, which leads to poor pre-labeling effects for categories with small amounts of data. Although some methods try to alleviate this problem through data augmentation or category weight adjustment, the effect is limited and the implementation is complex. In general, when dealing with category imbalance, the existing training pre-labeling process is prone to cause the model to be biased towards categories with large amounts of data, which in turn affects the pre-labeling effect and increases the labeling cost.

[0035] In view of this, some embodiments of the present application provide a data labeling method, which obtains training samples and a pre-trained multi-task labeling model, first trains a shared feature extraction network in the pre-trained multi-task labeling model using the training samples, and freezes its network parameters after the shared feature extraction network training is completed, and then uses the training samples to train the pre-trained multi-task labeling model with the network parameters of the shared feature extraction network frozen again, to obtain N trained classification networks, and then obtain the trained multi-task labeling model, so that only one model is trained to obtain a labeling model that can recognize multiple different categories of data. At the same time, the shared feature extraction network can be frozen after training, and there is no need to adjust the network parameters of the shared feature extraction network when subsequently training each classification network, which reduces the amount of data processing and can better maintain the consistency of the labeling results, thereby improving the efficiency of data labeling.

[0036] Figure 1 is a flow chart of a training method for a multi-task annotation model provided in an embodiment of the present application. Figure 1 As shown, the data labeling method includes the following steps:

[0037] In step S101, a training sample is obtained.

[0038] The training samples include an image data set, each image data in the image data set is respectively annotated with at least one category label, and the image data set corresponds to N different categories, where N is a positive integer greater than 1.

[0039] In step S102, a pre-trained multi-task annotation model is obtained.

[0040] Among them, the pre-trained multi-task annotation model includes a shared feature extraction network and N different classification networks, each classification network corresponds to a category.

[0041] In step S103, the pre-trained multi-task annotation model is trained using the training samples, and in response to determining that the shared feature extraction network has converged, the network parameters of the shared feature extraction network are frozen.

[0042] In step S104, the pre-trained multi-task labeling model with the network parameters of the shared feature extraction network frozen is trained again using the training samples until all N classification networks converge, thereby obtaining a trained multi-task labeling model.

[0043] In certain embodiments of the present application, the training method of the multi-task labeling model can be executed by a server, or by a terminal with certain computing capabilities. Further, the training method of the multi-task labeling model is used to train a multi-task labeling model that can label data in a data set with multiple different categories. Among them, the amount of data contained in some categories in the data set may be small, for example, less than a preset data amount threshold.

[0044] For example, a data set may be composed of multiple different categories of autonomous driving data, some of which, such as image data of trucks with protrusions or irregular vehicles, image data of vehicles making sharp turns, image data of pedestrians holding umbrellas, image data acquired under extreme weather or extremely poor lighting conditions, image data of trees falling in the middle of the road, etc., account for a very small proportion of the entire autonomous driving data set, but they belong to independent categories. These data can be called long-tail data.

[0045] In the related art, a general labeling model is usually trained to label all categories of data, or a separate labeling model is trained for different categories of data. The labeling accuracy of the former will inevitably be affected. When training the latter, if the data of this category is long-tail data, the training sample size priority will lead to insufficient accuracy of the labeling model obtained by training. Therefore, an embodiment of the present application proposes a multi-task labeling model training, which provides an efficient labeling model for autonomous driving data labeling by training a labeling model including a shared feature extraction network and N classification networks corresponding to different categories.

[0046] In some embodiments of the present application, a training sample may be first obtained, the training sample including an image data set, each image data in the image data set is respectively labeled with at least one category label, and the image data set corresponds to N different categories. In one example, the training sample may be an autonomous driving image data set, the image data set corresponds to multiple categories such as vehicles, people, roads, traffic lights, straight driving, turning, and stationary, that is, each image data in the image data set is labeled with at least one category label. Furthermore, some labels corresponding to the image data set may contain a small amount of data, such as the image data of the tree falling in the middle of the road, and its corresponding label may be, for example, a stationary obstacle across multiple roads.

[0047] On the other hand, a pre-trained multi-task annotation model can also be obtained, which includes a shared feature extraction network and N different classification networks, each classification network corresponding to a category. The shared feature extraction network can be located at the input layer of the annotation model to extract features from the input sample. The N classification networks can be located at the output layer of the annotation model to classify and annotate each image data in the input sample based at least on the features extracted by the shared feature extraction network.

[0048] In some implementations, the pre-trained multi-task labeling model can be trained using training samples, and when it is determined that the shared feature extraction network has converged, the network parameters of the shared feature extraction network are frozen. Next, the pre-trained multi-task labeling model with the network parameters of the shared feature extraction network frozen is trained again using the training samples until all N classification networks converge, thereby obtaining a trained multi-task labeling model.

[0049] According to the technical solution provided in the embodiment of the present application, by obtaining training samples and a pre-trained multi-task labeling model, the training samples are used to first train the shared feature extraction network in the pre-trained multi-task labeling model, and the network parameters of the shared feature extraction network are frozen after the training is completed. Then, the pre-trained multi-task labeling model with the network parameters of the shared feature extraction network frozen is trained again using the training samples to obtain N trained classification networks, and then the trained multi-task labeling model is obtained, so that only one model needs to be trained to obtain a labeling model that can recognize multiple different categories of data. At the same time, the shared feature extraction network can be parameter-frozen after training. When the classification networks are subsequently trained, there is no need to adjust the network parameters of the shared feature extraction network, which reduces the amount of data processing and can better maintain the consistency of the labeling results, thereby improving the efficiency of data labeling.

[0050] Figure 2 is a flow chart of a method for determining the convergence of a shared feature extraction network provided in an embodiment of the present application. Figure 2 As shown, the method comprises the following steps:

[0051] In step S201, the training samples are input into a pre-trained multi-task annotation model to obtain a first predicted classification of each image data in the image data set.

[0052] In step S202, the first predicted classification and the category label of each image data are compared, and in response to determining that the probability of incorrect labeling of each image data is greater than or equal to a first preset threshold, the network parameters of the shared feature extraction network are updated.

[0053] In step S203, the steps of inputting the training samples into the pre-trained multi-task labeling model to obtain the first predicted classification of each image data in the image data set, comparing the predicted classification and the category label of each image data, and updating the network parameters of the shared feature extraction network in response to determining that the error labeling probability of each image data is greater than or equal to the first preset threshold are repeatedly performed until it is determined that the error labeling probability of each image data is less than the first preset threshold.

[0054] In step S204, it is determined whether the shared feature extraction network converges.

[0055] In some embodiments of the present application, when training a multi-task annotation model, the pre-trained multi-task annotation model may be first trained using training samples until the shared feature extraction network converges. In one example, the training samples may be input into the pre-trained multi-task annotation model to obtain a first predicted classification of each image data in the image data set. At this point, the network parameters of the shared feature extraction network and the network parameters of each classification network are all initialization values.

[0056] After obtaining the first predicted classification of each image data in the image data set, the first predicted classification of each image data and the category label can be compared. If it is determined that the probability of incorrect labeling of each image data is greater than or equal to a first preset threshold, the network parameters of the shared feature extraction network are updated.

[0057] Further, the steps of inputting the training samples into the pre-trained multi-task annotation model to obtain the first predicted classification of each image data in the image data set, comparing the predicted classification and category label of each image data, and updating the network parameters of the shared feature extraction network in response to determining that the probability of error annotation of each image data is greater than or equal to the first preset threshold are repeatedly performed until it is determined that the probability of error annotation of each image data is less than the first preset threshold. The value of the first preset threshold is set according to actual needs and is not limited here. It is understandable that the first preset threshold can be set to a smaller value, that is, in this round of training, there is no need to require the model to achieve very high accuracy.

[0058] That is to say, when training the shared feature extraction network, only the network parameters of the shared feature extraction network can be iterated, instead of the network parameters of each classification network. Instead, a rough classification value is obtained based on the initialization network parameters of the classification network. If the rough classification value can meet the preset conditions, such as the probability of wrong labeling is less than the first preset threshold, it can be considered that the extraction of common features of various categories of data is relatively accurate. At this time, the shared feature extraction network can be considered to have converged, and the network parameters of the shared feature extraction network can be frozen. Next, the training data can be used again to train each classification network.

[0059] Figure 3 1 is a flow chart of a method for retraining a pre-trained multi-task annotation model with frozen network parameters of a shared feature extraction network using training samples provided in an embodiment of the present application until all N classification networks converge. Figure 3 As shown, the method comprises the following steps:

[0060] In step S301, the training samples are input into the pre-trained multi-task annotation model, and the common features of the training samples are extracted using a shared feature extraction network with frozen network parameters.

[0061] In step S302, the first target classification network is used to extract target-specific features of the training samples, and the common features and the target-specific features are used to obtain a second predicted classification of each image data in the image data set.

[0062] The first target classification network is any one of the N classification networks.

[0063] In step S303, the second predicted classification and category label of each image data are compared, and in response to determining that the probability of correct labeling of each image data is less than or equal to the second preset threshold, or the probability of incorrect labeling of each image data is greater than or equal to the third preset threshold, the network parameters of the first target classification network are updated.

[0064] In step S304, the steps of extracting target-specific features of training samples using the first target classification network, obtaining a second predicted classification of each image data in the image data set using the common features and the target-specific features, and updating the network parameters of the first target classification network in response to determining that the probability of correct labeling of each image data is less than or equal to the second preset threshold or the probability of incorrect labeling of each image data is greater than or equal to the third preset threshold are repeatedly performed until it is determined that the probability of correct labeling of each image data is greater than the second preset threshold and the probability of incorrect labeling of each image data is less than the third preset threshold.

[0065] In step S305, it is determined that the first target classification network converges.

[0066] In step S306 , in response to determining that all first target classification networks have converged, it is determined that all N classification networks have converged.

[0067] In some embodiments of the present application, when using training samples to train each classification network respectively, the training samples can be input into a pre-trained multi-task annotation model, and the common features of the training samples can be extracted using a shared feature extraction network with frozen network parameters. On the other hand, the first target classification network can also be used to extract target-specific features of the training samples. The use of the first target classification network to extract the target-specific features of the training samples can be that the first target classification network directly extracts the target-specific features from the training samples, or the first target classification network extracts the target-specific features from the common features.

[0068] In some implementations, the extracted common features and target-specific features can be used to obtain a second predicted classification of each image data in the image data set. The second predicted classification of each image data is compared with the category label. If it is determined that the probability of correct labeling of each image data is less than or equal to the second preset threshold, or the probability of incorrect labeling of each image data is greater than or equal to the third preset threshold, the network parameters of the first target classification network are updated. The values ​​of the second preset threshold and the third preset threshold are set according to actual needs and are not limited here.

[0069] That is to say, for the image data in the training sample that is labeled as the category corresponding to the first target classification network, its predicted classification value in the first target classification network should be higher than the first prediction threshold. Conversely, for the image data in the training sample that is labeled as the category corresponding to other classification networks except the first target classification network, its predicted classification value in the first target classification network should be lower than the second prediction threshold. The values ​​of the first prediction threshold and the second prediction threshold are set according to actual needs and are not limited here.

[0070] Next, the steps of using the first target classification network to extract target specific features of training samples, using common features and target specific features to obtain a second predicted classification of each image data in the image data set, and updating the network parameters of the first target classification network in response to determining that the probability of correct labeling of each image data is less than or equal to the second preset threshold or the probability of incorrect labeling of each image data is greater than or equal to the third preset threshold can be repeated until it is determined that the probability of correct labeling of each image data is greater than the second preset threshold and the probability of incorrect labeling of each image data is less than the third preset threshold.

[0071] If it is determined that the correct labeling probability of each image data is greater than the second preset threshold and the incorrect labeling probability of each image data is less than the third preset threshold, it can be further determined that the first target classification network converges. If it is determined that all first target classification networks converge, it can be determined that all N classification networks converge.

[0072] In some embodiments, the classification network includes at least a second target classification network, and the second target classification network is used to annotate long-tail data, where the long-tail data is data belonging to the target category whose amount of data is less than a preset data amount threshold. Furthermore, the second target classification network includes a fully connected layer and a classification layer, and also includes at least one regularization processing layer, and the learning rate of the second target classification network is less than a fourth preset threshold. The value of the fourth preset threshold is set according to actual needs and is not limited here.

[0073] In other words, you can use the labeled data to train a multi-category model that contains multiple categories, and use large amounts of category data to enhance the model's feature extraction capabilities. During training, you can first extract multi-category shared parameters from the pre-trained model. These parameters have a wide range of feature extraction capabilities and are applicable to multiple categories. Next, freeze the extracted shared parameters as a fixed feature extraction module to avoid updating these parameters in subsequent training.

[0074] Then configure an independent network head for each category and train them separately. Data enters the independent head through a fixed feature extraction module for category-specific optimization and learning. These independent heads are independent output layers designed and used for each task or output target in the multi-task annotation model, which can also be called heads. These heads are usually deployed after the shared layer of the multi-task annotation model to process different tasks through different output structures, and each head has its own parameters and structure, dedicated to a specific task.

[0075] For example, a multi-task annotation model can handle two tasks: image classification and object detection. The first few layers of the model can be used to extract common image features, and then two independent network heads are configured: the image classification head and the object detection head. The image classification head can be used to predict the category of the image, and the object detection head can be used to predict the location and category of the object in the image. In independent head training, the classification head and the detection head calculate their own loss functions and update parameters independently. This approach helps avoid negative interference between tasks and improve the performance of each task.

[0076] That is, the overall process of training a multi-task labeling model and then using the trained model to label the target data is to first collect and organize the labeled data set to ensure that the data set covers samples of various categories, some of which may have less data. Then select a model suitable for multi-category pre-labeling, such as a CNN (Convolutional Neural Networks) or Transformer model. Then train the multi-task labeling model. During training, first focus on the overall feature extraction ability of the model to ensure its generalization performance on multi-category data, evaluate the trained model to ensure that the training performance is normal, and then select shared parameters from the pre-trained model, usually the network backbone. The selected shared parameters are fixed and frozen to avoid updating them in subsequent independent head training to maintain their general feature extraction capabilities.

[0077] Next, we design an independent network head for each category, such as a detection head, a segmentation head, etc. The data of each category is input into the corresponding network head for training through a shared parameter module. During training, multiple network heads can be trained in parallel to improve training efficiency. When all network heads are trained, the trained multi-task labeling model can be obtained.

[0078] At this point, the trained multi-task labeling model can be used to label new data. The new data is first feature extracted through the shared parameter module, and then input into each independent head to obtain the pre-labeling results for each category. Combining the pre-labeling results of each independent head, the final labeled data can be generated. This approach can effectively alleviate the long-tail problem and improve the accuracy and efficiency of data pre-labeling.

[0079] In one example, if the training samples correspond to labels A1, A2, and A3, where labels A1 and A2 have more label sample data, while label A3 has less sample data. At this time, a pre-trained multi-task annotation model can be selected. The pre-trained multi-task annotation model is built based on the CNN model, and its input layer is used to receive input image data, such as the shape of the image input, including height, width, channels, etc. The input layer size is, for example, (224, 224, 3).

[0080] Its convolutional layer can include multiple groups of convolutional layers, each of which is followed by ReLU (Linear Rectification Function) for activation and batch normalization, such as Conv2D(32, kernel_size = (3, 3), activation = 'relu') and BatchNormalization(). Among them, Conv2D refers to a two-dimensional convolution operation, kernel_size is the convolution kernel size, activation = 'relu' triggers the ReLU function, and BatchNormalization() is the batch normalization function. Batch normalization can speed up training and stabilize the network.

[0081] The pooling layer of the model can be a maximum pooling layer, such as MaxPooling2D(pool_size=(2, 2)), where pool_size is the pooling size. The pooling layer is used to downsample and reduce the dimension of the feature map to retain important information. Alternatively, average pooling can also be used, which is not limited here.

[0082] The convolution layer can be used to extract local features and use convolution kernels (filters) for feature mapping.

[0083] The entire convolution and pooling part of the CNN model can be used as a shared feature extraction network to extract common features. The shared feature extraction network has shared parameters, which will not be updated during the independent head training phase.

[0084] Furthermore, an independent head can be designed for each category A1, A2 and A3. In one example, the head of A1 can be configured to include a fully connected layer, a random dropout layer and a classification layer, such as Dense(128, 'relu)+Dropout(0.5)+Dense(1, 'sigmoid'). Among them, Dense() is a function for creating a fully connected layer, and sigmoid is an S-shaped growth curve. The fully connected layer is used to connect to the classifier and output the final classification result. For example, the operation of the fully connected layer can be to flatten the feature map into a one-dimensional vector and perform a fully connected operation. The classification layer serves as an output layer to output the final prediction result, that is, the classification prediction result.

[0085] You can also configure the head of A2 to have the same structure as the head of A1, but with different weight initializations because they correspond to different categories. You can also configure the head of A3 to have the same structure as the heads of A1 and A2, but you may need more regularization or a smaller learning rate to accommodate smaller data. Alternatively, you can configure the heads of A1, A2, and A3 to have different structures, which is not limited here.

[0086] When training the model, you can use a large amount of data from A1, A2, and A3 to train the CNN model and extract the shared parameters. Then freeze the shared parameters, that is, lock the convolution and pooling layer parameters. Next, for independent head training, you can use a small amount of data from A1, A2, and A3 to train the independent heads respectively, while keeping the shared parameters unchanged. After the training is completed, you can use the trained CNN model to annotate the unlabeled image data.

[0087] In another example, if the training samples still correspond to labels A1, A2, and A3, where labels A1 and A2 have more label sample data, while label A3 has less sample data. At this time, a pre-trained multi-task annotation model can be selected. The pre-trained multi-task annotation model is built based on the Transformer model, which includes an input embedding layer for converting input data such as images into embedded vectors, that is, embedding into a high-dimensional vector space, such as through Embedding(input_dim=vocab_size,output_dim=embed_dim), the input image is expressed as a vector, where input_dim is the input dimension, vocab_size is the input atlas size, output_dim is the output dimension, and embed_dim is the dimension of the specified embedding vector.

[0088] Furthermore, the Transformer model also includes a position encoding layer for adding position information to the input so that the Transformer model can perceive the sequence order. The position encoding layer usually uses sine and cosine functions for position encoding.

[0089] The embedding layer, position encoding layer, and multiple encoder layers (each encoder layer contains a multi-head self-attention layer and a feedforward network) of the Transformer model constitute the shared feature extraction network of the model, which is used to extract common features. The shared feature extraction network has shared parameters, which will not be updated during the independent head training phase.

[0090] The multi-head self-attention layer is used to calculate the self-attention of the input sequence and generate multiple attention heads for each position. For example, MultiHeadAttention(num_heads=8, key_dim=embed_dim) where MultiHeadAttention is a multi-head attention mechanism, num_heads is the number of heads, and key_dim is the key vector dimension.

[0091] Feedforward networks are also called feedforward neural networks and are used for nonlinear transformation and feature extraction, such as Dense(units=2048, activation='relu'), Dense(units=embed_dim).

[0092] Furthermore, a separate head can be designed for each category A1, A2, and A3. In one example, the head of A1 can be configured to include a fully connected layer and a classification layer, such as Dense(128, 'relu)+Dense(1, 'sigmoid'). Dense() is a function for creating a fully connected layer, and sigmoid is an S-shaped growth curve.

[0093] In some implementations, the Transformer model may also include a normalization layer and a residual connection layer, which are potentially used to stabilize network training, such as LayerNormalization(). The latter retains input information through skip connections to avoid gradient disappearance, and its operation may be to add input to output.

[0094] The encoder of the Transformer model can be obtained by stacking multiple layers of self-attention layers and feed-forward networks to encode the input sequence. The decoder can be composed of self-attention layers, encoder-decoder attention layers and feed-forward networks to decode the target sequence and obtain the decoded output.

[0095] It is also possible to configure the head of A2 to have the same structure as the head of A1, but with different weight initializations because they correspond to different categories. It is also possible to configure the head of A3 to have the same structure as the heads of A1 and A2, but additional regularization measures may be required. Alternatively, it is also possible to configure the heads of A1, A2, and A3 to have different structures, which is not limited here.

[0096] When training the model, you can use a large amount of data from A1, A2, and A3 to train the Transformer model and extract the shared parameters. Then freeze the shared parameters, that is, lock the parameters of the embedding, position encoding, and encoder layers. Next, for independent head training, you can use a small amount of data from A1, A2, and A3 to train the independent heads respectively, while keeping the shared parameters unchanged. After the training is completed, you can use the trained CNN model to annotate the unlabeled image data.

[0097] In this way, training each independent head for each category separately can ensure that each category is optimized in a targeted manner, thereby helping the labeling model maintain good generalization capabilities when dealing with long-tail problems.

[0098] In some embodiments of the present application, if it is determined that M image data subsets corresponding to different new categories are added to the training sample, M classification networks are added to the multi-task annotation model, and the M classification networks correspond to the M new categories respectively, where M is a positive integer. Then, the multi-task annotation model is retrained using the updated training sample.

[0099] That is to say, if new image data is added to the training samples and the image data belongs to a new category, the performance of the model may decrease because the multi-task labeling model does not contain a classification network for classifying and predicting the new category. In this case, the multi-task labeling model can be retrained using the updated training samples.

[0100] Figure 4 is a flow chart of a method for retraining a multi-task annotation model using updated training samples provided in an embodiment of the present application. Figure 4 As shown, the method comprises the following steps:

[0101] In step S401, in response to determining that the image data in the M image data subsets are all long-tail data, the newly added M classification networks are trained using the updated training samples until the M classification networks converge to obtain a trained multi-task labeling model.

[0102] Among them, long-tail data refers to data belonging to the target category whose data volume is less than a preset data volume threshold.

[0103] In step S402, in response to determining that the image data in the M image data subsets includes at least one non-long-tail data, the multi-task annotation model is trained using the updated training samples, and in response to determining that the shared feature extraction network converges, the network parameters of the shared feature extraction network are frozen.

[0104] In step S403, the multi-task labeling model with the network parameters of the shared feature extraction network frozen is trained again using the training samples until all N+M classification networks converge, thereby obtaining a trained multi-task labeling model.

[0105] In some embodiments, when retraining the multi-task annotation model using the updated training samples, it can be first determined whether the image data in the M image data subsets include non-long-tail data. If not, that is, the image data in the M image data subsets are all long-tail data, then the updated training samples can be used to train the newly added M classification networks until the M classification networks converge to obtain the trained multi-task annotation model, and the long-tail data is the data whose amount of data belonging to the target category is less than the preset data amount threshold. At this time, there is no need to retrain the shared feature extraction network.

[0106] On the contrary, if the image data in the M image data subsets includes at least one non-long-tail data, the multi-task labeling model is trained using the updated training samples, and in response to determining that the shared feature extraction network has converged, the network parameters of the shared feature extraction network are frozen. The multi-task labeling model with the frozen network parameters of the shared feature extraction network is trained again using the training samples until all N+M classification networks converge to obtain the trained multi-task labeling model. That is, at this time, it is necessary to retrain the shared feature extraction network using the updated training samples, and freeze the network parameters of the shared feature extraction network again after training. Next, each classification network is trained using the updated training samples.

[0107] Figure 5 Schematic diagram of a data labeling method provided in an embodiment of the present application. Figure 5 As shown, the method comprises the following steps:

[0108] In step S501, target data is acquired.

[0109] In step S502, a multi-task annotation model is obtained.

[0110] Among them, the multi-task annotation model adopts Figures 1 to 4 A multi-task annotation model trained by the method provided by any of the embodiments shown.

[0111] In step S503, the trained multi-task labeling model is used to predict the target data to obtain a predicted classification.

[0112] In step S504, the target data is labeled using the predicted classification.

[0113] In certain embodiments of the present application, when using a trained multi-task annotation model to annotate unlabeled new data, the target data can be first obtained, and the multi-task annotation model trained in the above manner can be obtained. The trained multi-task annotation model is used to predict the target data to obtain a predicted classification. Finally, the target data is annotated using the predicted classification.

[0114] By adopting the technical solution provided in the embodiment of the present application, the labeling effect of a small number of data categories is enhanced by using a shared feature extraction network, and the accuracy of labeling is improved. At the same time, freezing the network parameters of the shared feature extraction network reduces the number of parameter updates during the training process, reduces the demand for computing resources, improves training efficiency, and ensures that each category is labeled in the same feature space, thereby improving the consistency of labeling between different categories. In addition, the design of the independent head enables each category to be optimized in a targeted manner, alleviating the imbalance caused by the long tail problem.

[0115] All the above optional technical solutions can be arbitrarily combined to form optional embodiments of the present application, which will not be described one by one here.

[0116] The following is an embodiment of the device of the present application, which can be used to execute the embodiment of the method of the present application. For details not disclosed in the embodiment of the device of the present application, please refer to the embodiment of the method of the present application.

[0117] Figure 6 Schematic diagram of a training device for a multi-task annotation model provided in an embodiment of the present application. Figure 6 As shown, the device comprises:

[0118] The acquisition module 601 is configured to acquire training samples, where the training samples include an image data set, each image data in the image data set is respectively annotated with at least one category label, and the image data set corresponds to N different categories, where N is a positive integer greater than 1.

[0119] The acquisition module 601 is further configured to acquire a pre-trained multi-task annotation model, where the pre-trained multi-task annotation model includes a shared feature extraction network and N different classification networks, each classification network corresponding to a category.

[0120] The training module 602 is configured to train the pre-trained multi-task annotation model using the training samples, and freeze the network parameters of the shared feature extraction network in response to determining that the shared feature extraction network has converged.

[0121] The training module 602 is further configured to use the training samples to retrain the pre-trained multi-task labeling model with the network parameters of the shared feature extraction network frozen until all N classification networks converge to obtain a trained multi-task labeling model.

[0122] According to the technical solution provided in the embodiment of the present application, by obtaining training samples and a pre-trained multi-task labeling model, the training samples are used to first train the shared feature extraction network in the pre-trained multi-task labeling model, and the network parameters of the shared feature extraction network are frozen after the training is completed. Then, the pre-trained multi-task labeling model with the network parameters of the shared feature extraction network frozen is trained again using the training samples to obtain N trained classification networks, and then the trained multi-task labeling model is obtained, so that only one model needs to be trained to obtain a labeling model that can recognize multiple different categories of data. At the same time, the shared feature extraction network can be parameter-frozen after training. When the classification networks are subsequently trained, there is no need to adjust the network parameters of the shared feature extraction network, which reduces the amount of data processing and can better maintain the consistency of the labeling results, thereby improving the efficiency of data labeling.

[0123] In some embodiments, determining that the shared feature extraction network has converged includes: inputting training samples into a pre-trained multi-task annotation model to obtain a first predicted classification for each image data in the image data set; comparing the first predicted classification and the category label of each image data, and in response to determining that the probability of an error labeling of each image data is greater than or equal to a first preset threshold, updating the network parameters of the shared feature extraction network; repeatedly executing the steps of inputting training samples into a pre-trained multi-task annotation model to obtain the first predicted classification for each image data in the image data set, comparing the predicted classification and the category label of each image data, and updating the network parameters of the shared feature extraction network in response to determining that the probability of an error labeling of each image data is greater than or equal to the first preset threshold, until determining that the probability of an error labeling of each image data is less than the first preset threshold; determining that the shared feature extraction network has converged.

[0124] In some embodiments, a pre-trained multi-task labeling model with frozen network parameters of a shared feature extraction network is trained again using training samples until all N classification networks converge, including: inputting the training samples into the pre-trained multi-task labeling model, and extracting common features of the training samples using the shared feature extraction network with frozen network parameters; extracting target-specific features of the training samples using a first target classification network, and obtaining a second predicted classification of each image data in an image data set using the common features and the target-specific features, wherein the first target classification network is any one of the N classification networks; comparing the second predicted classification of each image data with the category label, and in response to determining that the probability of correct labeling of each image data is less than or equal to a second preset threshold, or that the error label of each image data is less than or equal to a second preset threshold, When the probability of correct labeling of each image data is greater than or equal to the second preset threshold, the network parameters of the first target classification network are updated; the steps of extracting target specific features of the training samples using the first target classification network, obtaining a second predicted classification of each image data in the image data set using the common features and the target specific features, and updating the network parameters of the first target classification network in response to determining that the probability of correct labeling of each image data is less than or equal to the second preset threshold or the probability of incorrect labeling of each image data is greater than or equal to the third preset threshold, until it is determined that the probability of correct labeling of each image data is greater than the second preset threshold and the probability of incorrect labeling of each image data is less than the third preset threshold; determining that the first target classification network converges; and in response to determining that all the first target classification networks converge, determining that all N classification networks converge.

[0125] In some embodiments, the classification network includes at least a second target classification network, which is used to label long-tail data, where the long-tail data is data belonging to the target category and the amount of data is less than a preset data amount threshold; the second target classification network includes a fully connected layer and a classification layer, and also includes at least one regularization processing layer, and the learning rate of the second target classification network is less than a fourth preset threshold.

[0126] In some implementations, the method further includes: in response to determining that M image data subsets corresponding to different new categories are added to the training samples, M classification networks are added to the multi-task labeling model, the M classification networks correspond to the M new categories respectively, and M is a positive integer; and the multi-task labeling model is retrained using the updated training samples.

[0127] In some embodiments, the multi-task labeling model is retrained using updated training samples, including: in response to determining that the image data in the M image data subsets are all long-tail data, using the updated training samples to train the newly added M classification networks until the M classification networks are all converged to obtain the trained multi-task labeling model, where the long-tail data is data whose amount of data belonging to the target category is less than a preset data amount threshold; or in response to determining that the image data in the M image data subsets include at least one item of non-long-tail data, using the updated training samples to train the multi-task labeling model, and in response to determining that the shared feature extraction network converges, freezing the network parameters of the shared feature extraction network; using the training samples to train the multi-task labeling model with the network parameters of the shared feature extraction network frozen again until the N+M classification networks are all converged to obtain the trained multi-task labeling model.

[0128] In some implementations, the method further includes: obtaining target data; obtaining a trained multi-task labeling model; using the trained multi-task labeling model to predict the target data to obtain a predicted classification; and labeling the target data using the predicted classification.

[0129] It should be understood that the size of the serial numbers of the steps in the above embodiments does not mean the order of execution. The execution order of each process should be determined by its function and internal logic, and should not constitute any limitation on the implementation process of the embodiments of the present application.

[0130] Figure 7 Schematic diagram of an electronic device provided in an embodiment of the present application. Figure 7 As shown, the electronic device 7 of this embodiment includes: a processor 701, a memory 702, and a computer program 703 stored in the memory 702 and executable on the processor 701. When the processor 701 executes the computer program 703, the steps in the above-mentioned method embodiments are implemented. Alternatively, when the processor 701 executes the computer program 703, the functions of the modules / units in the above-mentioned device embodiments are implemented.

[0131] The electronic device 7 may be a desktop computer, a notebook, a PDA, a cloud server, or other electronic device. The electronic device 7 may include, but is not limited to, a processor 701 and a memory 702. Those skilled in the art will appreciate that Figure 7 The electronic device 7 is merely an example and does not limit the electronic device 7 , and may include more or less components than those shown in the figure, or different components.

[0132] The processor 701 may be a central processing unit (CPU), or other general-purpose processors, digital signal processors (DSP), application-specific integrated circuits (ASIC), field-programmable gate arrays (FPGA), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc.

[0133] The memory 702 may be an internal storage unit of the electronic device 7, for example, a hard disk or memory of the electronic device 7. The memory 702 may also be an external storage device of the electronic device 7, for example, a plug-in hard disk, a smart media card (SMC), a secure digital (SD) card, a flash card, etc. equipped on the electronic device 7. The memory 702 may also include both an internal storage unit of the electronic device 7 and an external storage device. The memory 702 is used to store computer programs and other programs and data required by the electronic device.

[0134] Those skilled in the art can clearly understand that for the convenience and simplicity of description, only the division of the above-mentioned functional units and modules is used as an example. In actual applications, the above-mentioned functions can be assigned to different functional units and modules as needed, that is, the internal structure of the device can be divided into different functional units or modules to complete all or part of the functions described above. The functional units and modules in the embodiments can be integrated into one processing unit, or each unit can exist physically separately, or two or more units can be integrated into one unit. The above-mentioned integrated unit can be implemented in the form of hardware or in the form of software functional units.

[0135] If the integrated module / unit is implemented in the form of a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, the present application implements all or part of the processes in the above-mentioned embodiment method, and can also be completed by instructing the relevant hardware through a computer program. The computer program can be stored in a computer-readable storage medium, and the computer program can implement the steps of the above-mentioned various method embodiments when executed by the processor. The computer program may include computer program code, which may be in source code form, object code form, executable file or some intermediate form. Computer-readable media may include: any entity or device capable of carrying computer program code, recording medium, U disk, mobile hard disk, disk, optical disk, computer memory, read-only memory (ROM), random access memory (RAM), electric carrier signal, telecommunication signal and software distribution medium, etc.

[0136] The above embodiments are only used to illustrate the technical solutions of the present application, rather than to limit them. Although the present application has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein. These modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the embodiments of the present application, and should all be included in the protection scope of the present application.

Claims

1. A training method for a multi-task annotation model, characterized in that: include: Acquire a training sample, wherein the training sample includes an image data set, each image data in the image data set is respectively marked with at least one category label, and the image data set corresponds to N different categories, where N is a positive integer greater than 1; Obtain a pre-trained multi-task annotation model, wherein the pre-trained multi-task annotation model includes a shared feature extraction network and N different classification networks, each classification network corresponding to a category; training the pre-trained multi-task annotation model using the training samples, and freezing network parameters of the shared feature extraction network in response to determining that the shared feature extraction network has converged; The pre-trained multi-task labeling model with the network parameters of the shared feature extraction network frozen is trained again using the training samples until all the N classification networks converge, thereby obtaining a trained multi-task labeling model.

2. The method according to claim 1, characterized in that The determining that the shared feature extraction network converges includes: Inputting the training samples into the pre-trained multi-task annotation model to obtain a first predicted classification of each image data in the image data set; Comparing the first predicted classification and the category label of each image data, and in response to determining that the probability of incorrect labeling of each image data is greater than or equal to a first preset threshold, updating the network parameters of the shared feature extraction network; Repeating the steps of inputting the training sample into the pre-trained multi-task annotation model to obtain a first predicted classification of each image data in the image data set, comparing the predicted classification and the category label of each image data, and updating the network parameters of the shared feature extraction network in response to determining that the probability of an incorrect annotation of each image data is greater than or equal to a first preset threshold, until determining that the probability of an incorrect annotation of each image data is less than the first preset threshold; Convergence of the shared feature extraction network is determined.

3. The method according to claim 1, characterized in that The method of retraining the pre-trained multi-task annotation model with frozen network parameters of the shared feature extraction network using the training samples until all the N classification networks converge includes: Inputting the training samples into a pre-trained multi-task annotation model, and using a shared feature extraction network with frozen network parameters to extract common features of the training samples; Extract target-specific features of the training samples using a first target classification network, and obtain a second predicted classification of each image data in the image data set using the common features and the target-specific features, wherein the first target classification network is any one of the N classification networks; Comparing the second predicted classification and the category label of each image data, and in response to determining that the correct labeling probability of each image data is less than or equal to the second preset threshold, or the incorrect labeling probability of each image data is greater than or equal to the third preset threshold, updating the network parameters of the first target classification network; Repeating the steps of extracting target-specific features of training samples using the first target classification network, obtaining a second predicted classification of each image data in the image data set using the common features and the target-specific features, and updating a network parameter of the first target classification network in response to determining that the probability of correct labeling of each image data is less than or equal to a second preset threshold or the probability of incorrect labeling of each image data is greater than or equal to a third preset threshold, until determining that the probability of correct labeling of each image data is greater than the second preset threshold and the probability of incorrect labeling of each image data is less than the third preset threshold; Determining that the first target classification network converges; In response to determining that all of the first target classification networks have converged, it is determined that the N classification networks have all converged.

4. The method according to claim 1, characterized in that: The classification network at least includes a second target classification network, and the second target classification network is used to label long-tail data, where the long-tail data is data whose amount of data belonging to the target category is less than a preset data amount threshold; The second target classification network includes a fully connected layer and a classification layer, and also includes at least one regularization processing layer, and the learning rate of the second target classification network is less than a fourth preset threshold.

5. The method according to claim 1, characterized in that The method further comprises: In response to determining that M image data subsets corresponding to different new categories are newly added to the training sample, M classification networks are newly added to the multi-task annotation model, and the M classification networks correspond to the M new categories respectively, where M is a positive integer; The multi-task labeling model is retrained using the updated training samples.

6. The method according to claim 5, characterized in that The retraining of the multi-task labeling model using the updated training samples includes: In response to determining that the image data in the M image data subsets are all long-tail data, using the updated training samples to train the newly added M classification networks until the M classification networks converge to obtain a trained multi-task labeling model, wherein the long-tail data is data belonging to a target category whose amount of data is less than a preset data amount threshold; or In response to determining that the image data in the M image data subsets includes at least one non-long-tail data, training the multi-task annotation model using the updated training sample, and in response to determining that the shared feature extraction network converges, freezing the network parameters of the shared feature extraction network; The multi-task labeling model with the network parameters of the shared feature extraction network frozen is trained again using the training samples until the N+M classification networks are all converged, thereby obtaining a trained multi-task labeling model.

7. The method according to claim 1, characterized in that After obtaining the trained multi-task labeling model, the method further includes: Get target data; Obtaining the trained multi-task labeling model; Using the trained multi-task annotation model to predict the target data, to obtain a predicted classification; The target data is labeled using the predicted classification.

8. A training device for a multi-task annotation model, characterized in that: include: An acquisition module is configured to acquire a training sample, wherein the training sample includes an image data set, each image data in the image data set is respectively marked with at least one category label, and the image data set corresponds to N different categories, where N is a positive integer greater than 1; The acquisition module is further configured to acquire a pre-trained multi-task annotation model, wherein the pre-trained multi-task annotation model includes a shared feature extraction network and N different classification networks, each classification network corresponding to a category; a training module configured to train the pre-trained multi-task annotation model using the training samples, and in response to determining that the shared feature extraction network has converged, freeze network parameters of the shared feature extraction network; The training module is also configured to use the training samples to retrain the pre-trained multi-task labeling model with the network parameters of the shared feature extraction network frozen until the N classification networks converge to obtain a trained multi-task labeling model.

9. An electronic device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, characterized in that: When the processor executes the computer program, the steps of the method according to any one of claims 1 to 7 are implemented.

10. A computer-readable storage medium storing a computer program, characterized in that: When the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 7 are implemented.

Images