Multi-task model training method based on adaptive loss function weighting

Abstract

The present invention discloses a multi-task model training method based on adaptive loss function weighting, including: 1) obtaining images and labels of a data set, and constructing a multi-task model to be trained; 2) inputting images into the multi-task model, and calculating the loss values ​​and initial gradients of different tasks; 3) normalizing the initial gradient of the ith task, and calculating the expected gradient of the ith task; 4) calculating the gradient difference through the normalized initial gradient of the ith task and its expected gradient, and updating its task weight; 5) repeating steps 3) to 4), updating the weights of N task losses, obtaining weighted multi-task losses, and updating the parameters of the multi-task model through a back propagation algorithm, and obtaining a trained multi-task model. The present invention adaptively weights the weights of different task losses, alleviates the gradient conflicts between tasks, and improves the accuracy and efficiency of multi-task model training.

Description

Technical Field

[0001] The present invention relates to the technical field of deep learning, and in particular to a multi-task model training method based on adaptive loss function weighting. Background Art

[0002] Most deep learning-based methods train a dedicated network model to solve a task. This type of method is called single-task learning. This method often ignores the correlation between tasks during model training, making the model prone to overfitting to the training data and reducing the generalization ability of the model. In addition, a network model is trained separately for each task, which increases the number of model parameters and redundancy.

[0003] In contrast, multi-task learning methods reduce the number of parameters of the entire model by sharing the same network model and training multiple tasks at the same time. The advantages of multi-task learning methods are also reflected in the introduction of additional inductive biases to reduce the risk of model overfitting, and it can fully exploit the similarities in the task feature space by utilizing the correlation between tasks. Many studies have shown that multi-task learning methods can make the model perform better on a single task than single-task learning models.

[0004] The current shortcomings in the application and development of multi-task learning methods mainly lie in the instability of their training process. Unlike single-task models, in the training of multi-task learning models, it is necessary to balance the weights of multiple parallel tasks in the overall loss function to prevent a certain task from dominating the training of the multi-task model. Previous training methods for multi-task models required tedious repeated experiments to perform grid searches on the weights of each subtask in the entire multi-task loss function, hoping to find the most appropriate weight. However, recent studies have shown that this weight search method is not only time-consuming, but also the search results are suboptimal. In fact, the optimal weight of each subtask changes dynamically during network training, and the fixed artificial weight assignment method cannot meet the training needs of multi-task models.

[0005] Based on the above discussion, the multi-task model training method based on adaptive loss function weighting has high practical application value. Summary of the invention

[0006] The purpose of the present invention is to overcome the shortcomings of the prior art and propose a multi-task model training method based on adaptive loss function weighting, which mainly corrects the gradients of loss items of different tasks, comprehensively considers the training rates of different tasks and the consistency of gradient update directions, adaptively adjusts the weights of different task losses in the overall loss, alleviates gradient conflicts between tasks, and prevents a certain task from dominating the training, thereby improving the accuracy and efficiency of multi-task model training.

[0007] To achieve the above object, the technical solution provided by the present invention is: a multi-task model training method based on adaptive loss function weighting, comprising the following steps:

[0008] 1) Obtain images and labels of the data set and build a multi-task model to be trained;

[0009] 2) Input the image into the multi-task model, record it as the tth iteration, initialize the weights of the N task losses, and calculate the loss values ​​and initial gradients of the N tasks at the tth iteration;

[0010] 3) Normalize the initial gradient of the i-th task and calculate the expected gradient of the i-th task;

[0011] 4) Calculate the gradient difference of the i-th task through the normalized initial gradient and expected gradient of the i-th task, and update the weight of the i-th task loss;

[0012] 5) Repeat steps 3) to 4), update the weights of the N task losses, iterate the multi-task model for the t+1th time, obtain the weighted multi-task loss, and update the parameters of the multi-task model through the back propagation algorithm. After multiple iterations, the trained multi-task model is obtained.

[0013] Further, in step 1), read the image in the dataset and the corresponding task label value Construct a multi-task model to be trained, which includes a feature extraction network, a feature fusion network, and a decoder network corresponding to different tasks.

[0014] Further, the step 2) comprises the following steps:

[0015] 2.1) Input the image, initialize the task loss weight and calculate the loss value of N tasks;

[0016] Input the image into the multi-task model and obtain the model's predicted value for the i-th task through the forward calculation of the model Initialize the loss weights ω of different tasks at the tth training iteration i (t0 = 1.0, and according to the loss function form of N tasks, calculate the difference between the label and the predicted value to obtain the loss term L of different tasks i (θ(t)) is expressed as:

[0017]

[0018] Where θ(t) represents the multi-task model parameters at the tth training iteration; Represents the loss function form of the i-th subtask, Represents the label value of the i-th task;

[0019] At the tth training iteration, the overall loss function L of the multi-task model MTL (θ(t)) is expressed as:

[0020]

[0021] In the formula, Represents the weighted sum of the losses of N tasks, where ω i (t) represents the weight of the i-th task loss at the t-th training iteration;

[0022] 2.2) Calculate the initial gradients of N tasks;

[0023] The loss terms of N tasks are passed through the back-propagation algorithm to obtain the gradient of the task-specific loss term, the initial gradient of the i-th task It is expressed as:

[0024]

[0025] In the formula, Represents ωi i (t)L i The gradient of (θ(t)) with respect to θ(t).

[0026] Further, the step 3) comprises the following steps:

[0027] 3.1) Calculate the normalized initial gradient of the i-th task;

[0028] After normalizing the initial gradient of the i-th task, it is expressed as:

[0029]

[0030] In the formula, represents the normalized initial gradient of the i-th task loss term, represents the initial gradient of the i-th task, log(L i (θ(t))+ε) refers to the i (θ(t))+ε is used to perform logarithmic operation, L i (θ(t)) represents the loss term for different tasks, and ε is a small amount used to increase the stability of the method;

[0031] 3.2) Calculate the normalized average loss reduction rate of the i-th task;

[0032] Get the loss value L of the i-th task calculated by the multi-task model in the first iteration i (θ(1)), The loss reduction rate calculated for the i-th task at the t-th iteration is calculated by the following formula:

[0033]

[0034] In the formula, t represents the number of current iteration training;

[0035] The average loss reduction rate R calculated by the i-th task at the t-th iteration i (t) is expressed as:

[0036]

[0037] In the formula, N represents the number of tasks, Indicates the weighted loss reduction rate corresponding to N tasks at the tth iteration;

[0038] Finally, the normalized average loss reduction rate V is calculated by the following formula: i (t):

[0039]

[0040] 3.3) Calculate the gradient correlation between the i-th task and the remaining tasks;

[0041] The weighted gradients of the remaining tasks are obtained by the following formula It is expressed as:

[0042]

[0043] In the formula, ω j (t) represents the weight of the remaining j-th subtask, L j (θ(t)) represents the loss term of the remaining j-th subtask, Represents ω j (t)L j The gradient of (θ(t)) with respect to θ(t);

[0044] The inner product of the gradient of the i-th task and the rest of the tasks is calculated using the following formula: It is expressed as:

[0045]

[0046] In the formula, Representative request and The inner product of

[0047] Traverse all other tasks and calculate the average of the inner product of the gradient of the i-th task and the rest of the tasks It represents the gradient correlation between the i-th subtask and the rest of the tasks:

[0048]

[0049] In the formula, Indicates that the gradient correlation between the i-th task and the remaining N-1 tasks is weighted;

[0050] 3.4) Calculate the expected gradient of the i-th task;

[0051] Correlate the gradient of the i-th task with the rest of the tasks The average loss reduction rate V calculated from the normalized i-th task at the t-th iteration i (t) is multiplied to get the expected gradient corresponding to the i-th task It is expressed as:

[0052]

[0053] Further, the step 4) comprises the following steps:

[0054] 4.1) Calculate the gradient difference of the i-th task;

[0055] Design a gradient-balanced loss function L G (ωi i (t)) to calculate the gradient difference of the i-th task, expressed as:

[0056]

[0057] In the formula, represents the normalized initial gradient of the i-th task loss term, represents the expected gradient of the loss value of the i-th task, Representative calculation and L 1 norm;

[0058] 4.2) Update the weight of the i-th task loss;

[0059] Calculate the gradient of the gradient balance loss function with respect to the task weight, and update the weight ωi of the i-th task loss at the t+1th iteration by the following formula i (t+1):

[0060]

[0061] In the formula, is the gradient-balanced loss function relative to the i-th task weight ω i (t), λ is a hyperparameter.

[0062] Further, the step 5) comprises the following steps:

[0063] 5.1) Weighting the loss items of different tasks;

[0064] Perform the t+1th iteration of the multi-task model, calculate the loss value of the i-th task at the t+1th iteration, and weight the loss items of different tasks using the following formula to obtain the multi-task loss function L at the t+1th iteration: MTL (θ(t+1)):

[0065]

[0066] In the formula, ω i (t+1), L i (θ(t+1)) are the weight and loss value of the i-th task at the t+1-th iteration;

[0067] 5.2) Update the parameters of the multi-task model;

[0068] Calculate the multi-task loss function L at the t+1th iteration through the back-propagation algorithm MTL The gradient of (θ(t+1)) with respect to the model parameter θ(t) And the parameters of the multi-task model are updated through the following formula, and the trained multi-task model is obtained through multiple iterations;

[0069]

[0070] Where θ(t+1) is the parameter of the multi-task model for the t+1th iteration, and λ is a hyperparameter.

[0071] Compared with the prior art, the present invention has the following advantages and beneficial effects:

[0072] 1. The present invention uses an adaptive loss function weighting method to automatically search and adjust the weights of different tasks. Compared with the method of manually searching for task weights, the search efficiency is higher, the speed is faster, and the effect is better.

[0073] 2. The multi-task models applicable to the present invention include not only existing classification models, but also target detection models, other complex semantic segmentation models and instance segmentation models, and other models using supervised learning.

[0074] 3. The present invention can flexibly adjust the number of algorithm iterations according to task requirements, and is easy and quick to use.

[0075] 4. Compared with other task weighted algorithms, the present invention has less calculation amount and faster operation speed, basically does not increase the time of training the network, and can further improve the training accuracy of the multi-task model. BRIEF DESCRIPTION OF THE DRAWINGS

[0076] Figure 1 The figure is a logic flow chart of the method of the present invention.

[0077] Figure 2 Schematic diagram of input images and corresponding labels for Task 1.

[0078] Figure 3 Schematic diagram of input images and corresponding labels for Task 2.

[0079] Figure 4 This is the prediction result diagram of the model output for task 1.

[0080] Figure 5 This is the prediction result diagram of the model output for task 2. DETAILED DESCRIPTION

[0081] The present invention is further described in detail below in conjunction with embodiments and drawings, but the embodiments of the present invention are not limited thereto.

[0082] like Figure 1 As shown, the multi-task model training method based on adaptive loss function weighting provided in this embodiment includes the following steps:

[0083] 1) The dataset uses the DTU fan defect dataset and reads the images and corresponding labels in the dataset. like Figure 2 , 3 As shown, Figure 2 The label box coordinates of task 1 are marked Figure 3 The label box coordinates of Task 2 are marked Construct a multi-task model to be trained, which includes a feature extraction network, a feature fusion network, and a decoder network corresponding to different tasks.

[0084] 2) Input the image into the multi-task model, record it as the tth iteration, initialize the weights of the N task losses, at this time N = 2, calculate the loss value and initial gradient of the N tasks at the tth iteration, including the following steps:

[0085] 2.1) Input the image, initialize the task loss weight and calculate the loss value of N tasks;

[0086] Input the image into the multi-task model and obtain the model's predicted value for the i-th task through the forward calculation of the model like Figure 4 and Figure 5 As shown, the predicted box coordinates of task 1 at this time The predicted box coordinates of task 2 are (2980, 1080, 2160, 1580) Initialize the loss weights ω of different tasks to (3290, 2120, 680, 290). i(t) = 1.0, and according to the loss function form of N tasks, calculate the difference between the label and the predicted value to obtain the loss term L for different tasks i (θ(t)) is expressed as:

[0087]

[0088] Where θ(t) represents the multi-task model parameters at the tth training iteration; Represents the loss function form of the i-th subtask, Represents the label value of the input image.

[0089] For example, using the predicted values ​​and flag values ​​of Task 1 and Task 2, the loss terms of the two tasks can be calculated as:

[0090]

[0091]

[0092] At the tth training iteration, the overall loss function L of the multi-task model MTL (θ(t)) is expressed as:

[0093]

[0094] In the formula, Represents the weighted sum of the losses of N tasks, where ω i (t) represents the weight of the i-th task loss at the t-th training iteration.

[0095] For example, using the predicted values ​​and flag values ​​of task 1 and task 2, the overall loss function L of the multi-task model is MTL (θ(t))=1.0×72+1.0×98=170.

[0096] 2.2) Calculate the initial gradients of N tasks;

[0097] The loss terms of N tasks are passed through the back-propagation algorithm to obtain the gradient of the task-specific loss term, the initial gradient of the i-th task It is expressed as:

[0098]

[0099] In the formula, Represents ω i (t)L i The gradient of (θ(t)) with respect to θ(t).

[0100] 3) Normalize the initial gradient of the i-th task and calculate the expected gradient of the i-th task, including the following steps:

[0101] 3.1) Calculate the normalized initial gradient of the i-th task, expressed as:

[0102]

[0103] In the formula, represents the normalized initial gradient of the i-th task loss term, represents the initial gradient of the i-th task, log(L i (θ(t))+ε) refers to the i (θ(t))+ε is used to perform logarithmic operation, L i (θ(t)) represents the loss term of different tasks, and ε is a small amount used to increase the stability of the method, which is set to 10 here. -6 .

[0104] 3.2) Calculate the normalized average loss reduction rate of the i-th task;

[0105] Get the loss value L of the i-th task calculated by the multi-task model in the first iteration i (θ(1)), The loss reduction rate calculated for the i-th task at the t-th iteration is calculated by the following formula:

[0106]

[0107] Where t represents the number of current iteration training.

[0108] The average loss reduction rate R calculated by the i-th task at the t-th iteration i (t), expressed as:

[0109]

[0110] In the formula, N represents the number of tasks, Indicates the weighted loss reduction rate corresponding to N tasks.

[0111] Finally, the normalized average loss reduction rate V is calculated by the following formula: i (t0:

[0112]

[0113] Where V i (t0 represents the normalized average loss reduction rate of the ith task calculated at the tth iteration, ε is a small amount used to increase the stability of the method, which is set to 10 here -6 .

[0114] 3.3) Calculate the gradient correlation between the i-th task and the remaining tasks;

[0115] The weighted gradients of the remaining tasks are obtained by the following formula It is expressed as:

[0116]

[0117] In the formula, ω j (t0 represents the weight of the remaining j-th subtask, L j (θ(t)) represents the loss term of the remaining j-th subtask, Represents ω j (t)L j The gradient of (θ(t)) with respect to θ(t).

[0118] The inner product of the gradient of the i-th task and the rest of the tasks is calculated using the following formula: It is expressed as:

[0119]

[0120] In the formula, Representative request and The inner product of .

[0121] Traverse all other tasks and calculate the average of the inner product of the gradient of the i-th task and the rest of the tasks It represents the gradient correlation between the i-th subtask and the rest of the tasks:

[0122]

[0123] In the formula, It represents the weighted correlation between the gradient of the i-th task and the remaining N-1 tasks.

[0124] 3.4) Calculate the expected gradient of the i-th task;

[0125] Correlate the gradient of the i-th task with the rest of the tasks The average loss reduction rate V calculated from the normalized i-th task at the t-th iteration i (t) is multiplied to get the expected gradient corresponding to the i-th task It is expressed as:

[0126]

[0127] 4) Calculate the gradient difference of the i-th task through the normalized initial gradient and expected gradient of the i-th task, and update the weight of the i-th task loss, including the following steps:

[0128] 4.1) Calculate the gradient difference of the i-th task;

[0129] Design a gradient-balanced loss function L G (ω i (t)) to calculate the gradient difference of the i-th task, expressed as:

[0130]

[0131] In the formula, represents the normalized initial gradient of the i-th task loss term, represents the expected gradient of the loss value of the i-th task, Representative calculation and L 1 Norm.

[0132] 4.2) Update the weight of the i-th task loss;

[0133] Calculate the gradient of the gradient balance loss function relative to the task weight. And update the weight ω of the i-th task loss at the t+1th iteration using the following formula: i (t+1):

[0134]

[0135] In the formula, is the gradient-balanced loss function relative to the i-th task weight ω i The gradient of (t), λ is a hyperparameter and is set to 0.001.

[0136] 5) Repeat steps 3) to 4) to update the weights of the N task losses. For example, the weights of task 1 and task 2 at the time t+1 are updated to ω1(t+1)=0.6 and ω2(t+1)=1.7 respectively. The multi-task model is iterated for the t+1th time to obtain the weighted multi-task loss, and the parameters of the multi-task model are updated through the back propagation algorithm. The trained multi-task model is obtained after multiple iterations, including the following steps:

[0137] 5.1) Weighting the loss items of different tasks;

[0138] Perform the t+1th iteration of the multi-task model, calculate the loss value of the i-th task at the t+1th iteration, and weight the loss items of different tasks using the following formula to obtain the multi-task loss function L at the t+1th iteration: MTL (θ(t+1)):

[0139]

[0140] In the formula, ω i (t+1), L i(t+1) are the weight and loss value of the i-th task at the t+1-th iteration respectively.

[0141] For example, using the updated weights of Task 1 and Task 2, calculate the multi-task loss function L for the t+1th iteration MTL (θ(t+1)) is:

[0142] L MTL (θ(t+1))=0.6L1(θ(t+1))+1.7L2(θ(t+1))

[0143] 5.2) Update the parameters of the multi-task model;

[0144] Calculate the multi-task loss function L at the t+1th iteration through the back-propagation algorithm MTL The gradient of (θ(t+1)) with respect to the model parameter θ(t) And the parameters of the multi-task model are updated through the following formula, and the trained multi-task model is obtained through multiple iterations.

[0145]

[0146] Where θ(t+1) is the parameter of the multi-task model for the t+1th iteration, and λ is a hyperparameter set to 0.001.

[0147] The above embodiments are preferred implementation modes of the present invention, but the implementation modes of the present invention are not limited to the above embodiments. Any other changes, modifications, substitutions, combinations, and simplifications that do not deviate from the spirit and principles of the present invention should be equivalent replacement methods and are included in the protection scope of the present invention.

Claims

1. Multi-task model training method based on adaptive loss function weighting, It is characterized in that The following steps are involved: 1) Obtain images and labels of the data set and build a multi-task model to be trained; 2) Input the image into the multi-task model, record it as the tth iteration, initialize the weights of the N task losses, and calculate the loss values ​​and initial gradients of the N tasks at the tth iteration, including the following steps: 2.1) Input the image, initialize the task loss weight and calculate the loss value of N tasks; Input the image into the multi-task model and obtain the model's predicted value for the i-th task through the forward calculation of the model Initialize the loss weights ω of different tasks at the tth training iteration i (t) = 1.0, and according to the loss function form of N tasks, calculate the difference between the label and the predicted value to obtain the loss term L for different tasks i (θ(t)) is expressed as: Where θ(t) represents the multi-task model parameters at the tth training iteration; Represents the loss function form of the i-th subtask, Represents the label value of the i-th task; At the tth training iteration, the overall loss function L of the multi-task model MTL (θ(t)) is expressed as: In the formula, Represents the weighted sum of the losses of N tasks, where ω i (t) represents the weight of the i-th task loss at the t-th training iteration; 2.2) Calculate the initial gradients of N tasks; The loss terms of N tasks are passed through the back-propagation algorithm to obtain the gradient of the task-specific loss term, the initial gradient of the i-th task It is expressed as: In the formula, Represents ω i (t)L i The gradient of (θ(t)) with respect to θ(t); 3) Normalize the initial gradient of the i-th task and calculate the expected gradient of the i-th task, including the following steps: 3.1) Calculate the normalized initial gradient of the i-th task; after normalizing the initial gradient of the i-th task, it is expressed as: In the formula, represents the normalized initial gradient of the i-th task loss term, represents the initial gradient of the i-th task, log(L i (θt))+ε) refers to L i (θ(t))+ε is used to perform logarithmic operation, L i (θ(t)) represents the loss term for different tasks, and ε is a small amount used to increase the stability of the method; 3.2) Calculate the normalized average loss reduction rate of the i-th task; Get the loss value L of the i-th task calculated by the multi-task model in the first iteration i (θ(1)), The loss reduction rate calculated for the i-th task at the t-th iteration is calculated using the following formula: In the formula, t represents the number of current iteration training; The average loss reduction rate R calculated by the i-th task at the t-th iteration i (t) is expressed as: In the formula, N represents the number of tasks, Indicates the weighted loss reduction rate corresponding to N tasks at the tth iteration; Finally, the normalized average loss reduction rate V is calculated by the following formula: i (t): 3.3) Calculate the gradient correlation between the i-th task and the remaining tasks; The weighted gradients of the remaining tasks are obtained by the following formula It is expressed as: In the formula, ω j (t) represents the weight of the remaining j-th subtask, L j (θ(t)) represents the loss term of the remaining j-th subtask, Represents ω j (t)L j The gradient of (θ(t)) with respect to θ(t); The inner product of the gradient of the i-th task and the rest of the tasks is calculated using the following formula: It is expressed as: In the formula, Representative request and The inner product of Traverse all other tasks and calculate the average of the inner product of the gradient of the i-th task and the rest of the tasks It represents the gradient correlation between the i-th subtask and the rest of the tasks: In the formula, Indicates that the gradient correlation between the i-th task and the remaining N-1 tasks is weighted; 3.4) Calculate the expected gradient of the i-th task; Correlate the gradient of the i-th task with the rest of the tasks The average loss reduction rate V calculated from the normalized i-th task at the t-th iteration i (t) is multiplied to get the expected gradient corresponding to the i-th task It is expressed as: 4) Calculate the gradient difference of the i-th task through the normalized initial gradient and expected gradient of the i-th task, and update the weight of the i-th task loss, including the following steps: 4.1) Calculate the gradient difference of the i-th task; Design a gradient-balanced loss function L G (ω i (t)) to calculate the gradient difference of the i-th task, expressed as: In the formula, represents the normalized initial gradient of the i-th task loss term, represents the expected gradient of the loss value of the i-th task, Representative calculation and L 1 norm; 4.2) Update the weight of the i-th task loss; Calculate the gradient of the gradient balance loss function with respect to the task weight, and update the weight ω of the i-th task loss at the t+1th iteration by the following formula i (t+1): In the formula, is the gradient-balanced loss function relative to the i-th task weight ω i The gradient of (t), λ is a hyperparameter; 5) Repeat steps 3) to 4), update the weights of the N task losses, iterate the multi-task model for the t+1th time, obtain the weighted multi-task loss, and update the parameters of the multi-task model through the back propagation algorithm. After multiple iterations, the trained multi-task model is obtained, including the following steps: 5.1) Weighting the loss items of different tasks; Perform the t+1th iteration of the multi-task model, calculate the loss value of the i-th task at the t+1th iteration, and weight the loss items of different tasks using the following formula to obtain the multi-task loss function L at the t+1th iteration: MTL (θ(t+1)): In the formula, ω i (t+1), L i (θ(t+1)) are the weight and loss value of the i-th task at the t+1-th iteration; 5.2) Update the parameters of the multi-task model; Calculate the multi-task loss function L at the t+1th iteration through the back-propagation algorithm MTL The gradient of (θ(t+1)) with respect to the model parameter θ(t) And the parameters of the multi-task model are updated through the following formula, and the trained multi-task model is obtained through multiple iterations; Where θ(t+1) is the parameter of the multi-task model for the t+1th iteration, and λ is a hyperparameter.

2. The multi-task model training method based on adaptive loss function weighting according to claim 1, It is characterized in that In step 1), read the image in the dataset and the corresponding task label value Construct a multi-task model to be trained, which includes a feature extraction network, a feature fusion network, and a decoder network corresponding to different tasks.

Images