Teaching student network for end-to-end semi-supervised object detection
A system and method for end-to-end semi-supervised object detection is provided. The system retrieves labeled and unlabeled images from an image dataset and generates an input batch by application of image transformation(s) on the images. The system further generates a first result for each image of the input batch by application of a teacher neural network on the input batch. For an object in an unlabeled image of the batch, the first result includes candidate bounding boxes and scores for the boxes. The system determines a threshold score based on the scores and selects a foreground bounding box from the candidates. The system generates a second result by application of a student neural network on the unlabeled image and computes a training loss over the input batch based on the foreground bounding box and the second result. The system trains the student neural network based on the training loss.
1. A method, comprising:
retrieving labeled and unlabeled images from an image dataset;
generating an input batch by an application of a set of image transformations on the labeled and unlabeled images;
generating a first result for each image of the input batch by an application of a teacher neural network on the input batch,
wherein the teacher neural network is a pretrained network for an object detection task, and
the first result for an object in a first unlabeled image of the input batch includes a set of candidate bounding boxes for the object and a set of scores corresponding to the set of candidate bounding boxes;
determining a threshold score based on the set of scores;
selecting, based on the threshold score, a foreground bounding box from the set of candidate bounding boxes;
generating a second result that includes a bounding box prediction for the object, by an application of a student neural network on the first unlabeled image,
wherein the student neural network is an untrained network that is to be trained for the object detection task;
computing a training loss over the input batch based on the foreground bounding box and the bounding box prediction; and
re-training the student neural network on the object detection task based on the training loss.
2. The method according to claim 1 , further comprising randomly sampling the image dataset using a sample ratio to retrieve the labeled and unlabeled images.
3. The method according to claim 1 , wherein the set of image transformations comprises a first subset of image transformations associated with a first data augmentation type and a second subset of image transformations associated with a second data augmentation type that is different from the first data augmentation type.
4. The method according to claim 3 , wherein the first subset of image transformations comprises an image flipping operation and an image shifting operation, and
the second subset of image transformations comprises one or more of an image rotation operation, a blur operation, a contrast variation, a shear operation, a masking operation for one or more regions of an image, a jitter addition operation, or an addition of a random noise.
5. The method according to claim 3 , wherein the input batch comprises a second unlabeled image and a labeled image associated with the first data augmentation type, and the first unlabeled image is associated with the second data augmentation type.
6. The method according to claim 1 , further comprising:
generating a third result by the application of the student neural network on a labeled image of the input batch; and
computing a total supervised loss for the first result associated with the labeled image and the third result associated with the labeled image, by using a supervised loss function and a supervised regression loss function,
wherein the total supervised loss includes a supervised classification loss and a supervised box-regression loss, and
the training loss is computed further based on the computation of the total supervised loss for the labeled image of the input batch.
7. The method according to claim 1 , further comprising:
computing a first unsupervised loss for the first result that is generated for the first unlabeled image;
computing a second unsupervised loss for the first result that is generated for a second unlabeled image of the input batch,
wherein each of the first unsupervised loss and the second unsupervised loss is computed by using an unsupervised loss function and includes an unsupervised classification loss and an unsupervised box-regression loss, and
wherein the training loss is computed further based on the computation of the first unsupervised loss and the second unsupervised loss.
8. The method according to claim 7 , wherein each of the first unsupervised loss and the second unsupervised loss is equal to a sum of a foreground classification loss, a background classification loss, a background similarity loss, and a foreground-background dissimilarity loss.
9. The method according to claim 1 , further comprising an application of a non-maximum suppression operation on the set of candidate bounding boxes to extract a subset of candidate bounding boxes from the set of candidate bounding boxes.
10. The method according to claim 9 , wherein the foreground bounding box is selected from the subset of candidate bounding boxes.
11. The method according to claim 9 , further comprising:
selecting a bounding box from the subset of candidate bounding boxes;
generate a set of jitter boxes by application of a jitter operation on the selected bounding box for a number of iterations; and
executing a bagging operation on the set of jitter boxes to select a jitter box with an area that is a maximum among respective areas of the set of jitter boxes,
wherein the selected jitter box is used in computing a box-regression loss that is part of the training loss.
12. The method according to claim 1 , wherein the set of scores includes:
foreground scores for foreground bounding boxes of the set of candidate bounding boxes, and
background scores for background bounding boxes of the set of candidate bounding boxes.
13. The method according to claim 12 , further comprising:
computing an average foreground score by dividing a sum of the foreground scores by a number of the foreground bounding boxes; and
computing an average background score by dividing a sum of the background scores by a number of the background bounding boxes, wherein the threshold score is determined by:
dividing the average foreground score by the average background score to produce a value; and
applying a floor function on the value.
14. The method according to claim 1 , further comprising updating weight parameters of the student neural network using the training loss to re-train the student neural network.
15. The method according to claim 14 , further comprising updating weight parameters of the teacher neural network based on the updated weight parameters of the student neural network.
16. The method according to claim 15 , wherein the update of the weight parameters of the teacher neural network comprises an execution of an exponential moving average (EMA) operation and an execution of an Exponential Adaptive Difference Moving Average (E-ADMA) operation.
17. A system, comprising:
circuitry configured to:
retrieve labeled and unlabeled images from an image dataset;
generate an input batch by an application of a set of image transformations on the labeled and unlabeled images;
generate a first result for each image of the input batch by an application of a teacher neural network on the input batch,
wherein the teacher neural network is a pretrained network for an object detection task, and
the first result for an object in a first unlabeled image of the input batch includes a set of candidate bounding boxes for the object and a set of scores corresponding to the set of candidate bounding boxes;
determine a threshold score based on the set of scores;
select, based on the threshold score, a foreground bounding box from the set of candidate bounding boxes;
generate a second result that includes a bounding box prediction for the object, by an application of a student neural network on the first unlabeled image,
wherein the student neural network is an untrained network that is to be trained for the object detection task;
compute a training loss over the input batch based on the foreground bounding box and the bounding box prediction; and
re-train the student neural network on the object detection task based on the training loss.
18. The system according to claim 17 , wherein the circuitry is further configured to update weight parameters of the student neural network using the training loss to re-train the student neural network.
19. The system according to claim 18 , wherein the circuitry is further configured to update weight parameters of the teacher neural network based on the updated weight parameters of the student neural network, and
the update of the weight parameters of the teacher neural network comprises an execution of an exponential moving average (EMA) operation and an execution of an Exponential Adaptive Difference Moving Average (E-ADMA) operation.
20. A non-transitory computer-readable storage medium configured to store instructions that, in response to being executed, causes a computer in a system to perform operations, the operations comprising:
retrieving labeled and unlabeled images from an image dataset;
generating an input batch by an application of a set of image transformations on the labeled and unlabeled images;
generating a first result for each image of the input batch by an application of a teacher neural network on the input batch,
wherein the teacher neural network is a pretrained network for an object detection task, and
the first result for an object in a first unlabeled image of the input batch includes a set of candidate bounding boxes for the object and a set of scores corresponding to the set of candidate bounding boxes;
determining a threshold score based on the set of scores;
selecting, based on the threshold score, a foreground bounding box from the set of candidate bounding boxes;
generating a second result that includes a bounding box prediction for the object, by an application of a student neural network on the first unlabeled image,
wherein the student neural network is an untrained network that is to be trained for the object detection task;
computing a training loss over the input batch based on the foreground bounding box and the bounding box prediction; and
re-training the student neural network on the object detection task based on the training loss.