Balancing multi-task learning through conditional or parallel batch mining
Methods, systems, and computer program products, for training a multitask neural network. First and second datasets are provided, containing samples for a first task and a second task, respectively. First and second batch miners are provided for mining samples from the first and second datasets. First and second costs for completing the first and second tasks, respectively, are assessed using a first sample mined by the first batch miner from the first dataset and a second sample mined by the second batch miner from the second dataset. When the first or second cost, respectively, falls within a range delimited by lower and upper thresholds, the is added to a first or second batch, respectively. When a termination condition is reached for either the first or second batch, the first or the second batch is used to update the neural network.
1. A computer-implemented method for training a multi-task neural network to detect features of an image captured by a camera, the method comprising:
providing at least a first dataset containing a first set of annotated images, and a second dataset containing a second set of annotated images;
providing a first batch miner for mining annotated images from the first dataset and a second batch miner for mining annotated images from the second dataset;
determining a first loss value for a first task, wherein the first task comprises processing, by the neural network, a first annotated image mined by the first batch miner from the first dataset;
determining a second loss value for a second task, wherein the second task comprises processing, by the neural network, a second annotated image mined by the second batch miner from the second data set;
if the first loss value falls within a range delimited by a first lower threshold value and a first upper threshold value, including the first annotated image in a first batch, wherein the first lower and upper threshold values are defined based on a current level for training the neural network;
if the second loss value falls within a range delimited by a second lower threshold value and a second upper threshold value, including the second annotated image in a second batch, wherein the second lower and upper threshold values are defined based on the current level for training the neural network; and
in response to reaching a termination condition for either the first batch or the second batch, using the first batch or the second batch to train the neural network.
2. The method of claim 1 , further comprising:
repeating the determining and including steps until the termination condition is reached; and
wherein the termination condition includes one or more of:
the first batch or the second batch reaching a pre-determined size, and
a pre-determined number of repetitions of the determining and including steps has been performed.
3. The method of claim 1 , wherein determining the first loss value and determining the second loss value are performed in sequence, one after another.
4. The method of claim 1 , wherein determining the first loss value and determining the second loss value are performed in parallel, and wherein the batch that is completed first is used to train the neural network.
5. The method of claim 1 , further comprising:
starting a new first batch after the first batch has been used to train the neural network; and
starting a new second batch after the second batch has been used to train the neural network.
6. The method of claim 1 , further comprising:
starting a new first batch and a new second batch after the first batch or the second batch has been used to train the neural network.
7. The method of claim 1 , wherein the first and second loss values are determined based on processing the annotated images by the neural network and comparing a result of the processing with respect to ground truths for the annotated images.
8. The method of claim 1 , wherein the first lower threshold value is the same as the second lower threshold value, and wherein the first upper threshold value is the same as the second upper threshold value.
9. The method of claim 1 , wherein the number of annotated images in the first batch is equal to the number of annotated images in the second batch.
10. A system for training a multitask neural network to detect features of an image captured by a camera, comprising:
a memory; and
a processor,
wherein the memory contains instructions that when executed by the processor causes the processor to perform a method that includes:
providing at least a first dataset containing a first set of annotated images and a second dataset containing a second set of annotated images;
providing a first batch miner for mining annotated images from the first dataset and a second batch miner for mining annotated images from the second dataset;
determining a first loss value for a first task, wherein the first task comprises processing, by the neural network, a first annotated image mined by the first batch miner from the first dataset;
determining a second loss value for a second task, wherein the second task comprises processing, by the neural network, a second annotated image mined by the second batch miner from the second data set;
if the first loss value falls within a range delimited by a first lower threshold value and a first upper threshold value, including the first annotate image in a first batch, wherein the first lower and upper threshold values are defined based on a current level for training the neural network;
if the second loss value falls within a range delimited by a second lower threshold value and a second upper threshold value, including the second annotate image in a second batch, wherein the second lower and upper threshold values are defined based on the current level for training the neural network; and
in response to reaching a termination condition for either the first batch or the second batch, using the first batch or the second batch to train the neural network.
11. A non-transitory computer-readable storage medium having stored thereon instructions being executable by a processor to perform a method comprising:
providing at least a first dataset containing a first set of annotated images and a second dataset containing a second set of annotated images;
providing a first batch miner for mining annotated images from the first dataset and a second batch miner for mining annotated images from the second dataset;
determining a first loss value for a first task, wherein the first task comprises processing, by a neural network, a first annotated image mined by the first batch miner from the first dataset;
determining a second loss value for a second task, wherein the second task comprises processing, by the neural network, a second annotated image mined by the second batch miner from the second data set;
if the first loss value falls within a range delimited by a first lower threshold value and a first upper threshold value, including the first annotated image in a first batch, wherein the first lower and upper threshold values are defined based on a current level for training the neural network;
if the second loss value falls within a range delimited by a second lower threshold loss value and a second upper threshold loss value, include the second annotated image in a second batch, wherein the second lower and upper threshold values are defined based on the current level for training the neural network; and
in response to reaching a termination condition for either the first batch or the second batch, using the first batch or the second batch to train the neural network.