Training distilled machine learning models using a pre-trained feature extractor
Methods, systems, and apparatus, including computer programs encoded on a computer storage medium, for training a student machine learning model using a teacher machine learning model that has a pre-trained feature extractor. In one aspect, a method includes obtaining data specifying the teacher machine learning model that is configured to perform a machine learning task; obtaining first training data; training the teacher machine learning model on the first training data to obtain a trained teacher machine learning model; generating second, automatically labeled training data by using the trained teacher machine learning model to process unlabeled training data; and training a student machine learning model to perform the machine learning task using at least the second, automatically labeled training data, wherein the student machine learning model does not include the pre-trained feature extractor and instead includes a different feature extractor having fewer parameters than the pre-trained feature extractor.
1 . A computer-implemented method comprising:
obtaining data specifying a pre-trained feature extractor that (i) is configured to receive an input image and to process the input image to generate an embedding of the input image and (ii) has been pre-trained as part of an image processing neural network that is configured through training to perform an image processing machine learning task by generating outputs for the image processing machine learning task from the embeddings generated by the pre-trained feature extractor, wherein the data specifying the pre-trained feature extractor comprises data specifying pre-trained values of parameters of the pre-trained feature extractor that have been determined as a result of the training of the image processing neural network;
generating, based on the data specifying the pre-trained feature extractor, a teacher machine learning model configured to perform a target machine learning task, wherein the target machine learning task is different from the image processing machine learning task that the image processing neural network is configured to perform, and wherein the teacher machine learning model comprises the pre-trained feature extractor having parameters that have the pre-trained values and a remaining portion having parameters that have initial values;
obtaining first training data comprising a plurality of first training inputs each associated with a ground truth output;
training the teacher machine learning model that comprises the pre-trained feature extractor and the remaining portion to perform the target machine learning task on the first training data to obtain a trained teacher machine learning model, wherein training the teacher machine learning model comprises updating values of parameters of the teacher machine learning model from the pre-trained values of the parameters of the pre-trained feature extractor and from the initial values of the parameters of the remaining portion, and wherein training the teacher machine learning model comprises holding the parameters of the pre-trained feature extractor fixed to the pre-trained values while updating the initial values of the parameters of the remaining portion;
generating second, automatically labeled training data by using the trained teacher machine learning model to process unlabeled training data comprising a plurality of second training inputs to generate, for each second training input, a pseudo ground truth output that is a target output that should be generated by a student machine learning model from processing the second training input; and
training the student machine learning model to perform the target machine learning task using at least the second, automatically labeled training data, wherein the student machine learning model does not include the pre-trained feature extractor included in the teacher machine learning model, and instead includes a different feature extractor having fewer parameters than the pre-trained feature extractor.
2 . The method of claim 1 , wherein the teacher machine learning model is configured to perform the target machine learning task by processing a teacher model input comprising the input image and point cloud data.
3 . The method of claim 1 , wherein training the teacher machine learning model on the first training data to obtain the trained teacher machine learning model comprises:
processing the first training input using the pre-trained feature extractor to generate an embedding of the first training input;
generating a teacher model training output for the target machine learning task from the embedding of the first training input; and
training the teacher machine learning model based on a difference between the teacher model training output and the ground truth output associated with the first training input.
4 . The method of claim 1 , wherein the pseudo ground truth output for each second training input is a soft target output.
5 . The method of claim 1 , wherein the pseudo ground truth output for each second training input is a hard target output.
6 . The method of claim 1 , wherein training the student machine learning model to perform the target machine learning task comprises fine-tuning the student machine learning model using the first training data after training the student machine learning model using the second, automatically labeled training data.
7 . The method of claim 1 , wherein the unlabeled training data has a larger volume than that of the first training data.
8 . The method of claim 1 , wherein generating the second, automatically labeled training data comprises applying data balancing techniques including class balancing techniques.
9 . The method of claim 8 , wherein the class balancing techniques comprise one or more of:
generating duplicated pseudo ground truth outputs for a single second training input,
randomly sampling a subset of the pseudo ground truth outputs from the pseudo ground truth outputs generated for the plurality of second training inputs, or
selecting the subset of the pseudo ground truth outputs according to values of the pseudo ground truth outputs.
10 . The method of claim 1 , further comprising deploying the student machine learning model at an on-board system of a vehicle.
11 . The method of claim 10 , wherein:
the target machine learning task comprises an image classification task; and
the student machine learning model is configured to perform the image classification task by processing input data generated by one or more camera sensors of the vehicle.
12 . The method of claim 1 , wherein the different machine learning task performed by the different image processing neural network comprises an image understanding task.
13 . A system comprising one or more computers and one or more storage devices storing instructions that are operable, when executed by the one or more computers, to cause the one or more computers to perform operations comprising:
obtaining data specifying a pre-trained feature extractor that (i) is configured to receive an input image and to process the input image to generate an embedding of the input image and (ii) has been pre-trained as part of an image processing neural network that is configured through training to perform an image processing machine learning task by generating outputs for the image processing machine learning task from the embeddings generated by the pre-trained feature extractor, wherein the data specifying the pre-trained feature extractor comprises data specifying pre-trained values of parameters of the pre-trained feature extractor that have been determined as a result of the training of the image processing neural network;
generating, based on the data specifying the pre-trained feature extractor, a teacher machine learning model configured to perform a target machine learning task, wherein the target machine learning task is different from the image processing machine learning task that the image processing neural network is configured to perform, and wherein the teacher machine learning model comprises the pre-trained feature extractor having parameters that have the pre-trained values and a remaining portion having parameters that have initial values;
obtaining first training data comprising a plurality of first training inputs each associated with a ground truth output;
training the teacher machine learning model that comprises the pre-trained feature extractor and the remaining portion to perform the target machine learning task on the first training data to obtain a trained teacher machine learning model, wherein training the teacher machine learning model comprises updating values of parameters of the teacher machine learning model from the pre-trained values of the parameters of the pre-trained feature extractor and from the initial values of the parameters of the remaining portion, and wherein training the teacher machine learning model comprises holding the parameters of the pre-trained feature extractor fixed to the pre-trained values while updating the initial values of the parameters of the remaining portion;
generating second, automatically labeled training data by using the trained teacher machine learning model to process unlabeled training data comprising a plurality of second training inputs to generate, for each second training input, a pseudo ground truth output that is a target output that should be generated by a student machine learning model from processing the second training input; and
training the student machine learning model to perform the target machine learning task using at least the second, automatically labeled training data, wherein the student machine learning model does not include the pre-trained feature extractor included in the teacher machine learning model, and instead includes a different feature extractor having fewer parameters than the pre-trained feature extractor.
14 . The system of claim 13 , wherein the teacher machine learning model is configured to perform the target machine learning task by processing a teacher model input comprising the input image and point cloud data.
15 . The system of claim 13 , wherein training the teacher machine learning model on the first training data to obtain the trained teacher machine learning model comprises:
processing the first training input using the pre-trained feature extractor to generate an embedding of the first training input;
generating a teacher model training output for the target machine learning task from the embedding of the first training input; and
determining an update to parameter values of the teacher machine learning model based on a difference between the teacher model training output and the ground truth output associated with the first training input.
16 . The system of claim 13 , wherein training the student machine learning model to perform the target machine learning task comprises fine-tuning the student machine learning model using the first training data after training the student machine learning model using the second, automatically labeled training data.
17 . The system of claim 13 , wherein the unlabeled training data has a larger volume than that of the first training data.
18 . The system of claim 13 , wherein generating the second, automatically labeled training data comprises applying data balancing techniques including class balancing techniques.
19 . The system of claim 18 , wherein the class balancing techniques comprise one or more of:
generating duplicated pseudo ground truth outputs for a single second training input,
randomly sampling a subset of the pseudo ground truth outputs from the pseudo ground truth outputs generated for the plurality of second training inputs, or
selecting the subset of the pseudo ground truth outputs according to values of the pseudo ground truth outputs.
20 . One or more computer storage media encoded with instructions that, when executed by one or more computers, cause the one or more computers to perform operations comprising:
obtaining data specifying a pre-trained feature extractor that (i) is configured to receive an input image and to process the input image to generate an embedding of the input image and (ii) has been pre-trained as part of an image processing neural network that is configured through training to perform an image processing machine learning task by generating outputs for the image processing machine learning task from the embeddings generated by the pre-trained feature extractor, wherein the data specifying the pre-trained feature extractor comprises data specifying pre-trained values of parameters of the pre-trained feature extractor that have been determined as a result of the training of the image processing neural network;
generating, based on the data specifying the pre-trained feature extractor, a teacher machine learning model configured to perform a target machine learning task, wherein the target machine learning task is different from the image processing machine learning task that the image processing neural network is configured to perform, and wherein the teacher machine learning model comprises the pre-trained feature extractor having parameters that have the pre-trained values and a remaining portion having parameters that have initial values;
obtaining first training data comprising a plurality of first training inputs each associated with a ground truth output;
training the teacher machine learning model that comprises the pre-trained feature extractor and the remaining portion to perform the target machine learning task on the first training data to obtain a trained teacher machine learning model, wherein training the teacher machine learning model comprises updating values of parameters of the teacher machine learning model from the pre-trained values of the parameters of the pre-trained feature extractor and from the initial values of the parameters of the remaining portion, and wherein training the teacher machine learning model comprises holding the parameters of the pre-trained feature extractor fixed to the pre-trained values while updating the initial values of the parameters of the remaining portion;
generating second, automatically labeled training data by using the trained teacher machine learning model to process unlabeled training data comprising a plurality of second training inputs to generate, for each second training input, a pseudo ground truth output that is a target output that should be generated by a student machine learning model from processing the second training input; and
training the student machine learning model to perform the target machine learning task using at least the second, automatically labeled training data, wherein the student machine learning model does not include the pre-trained feature extractor included in the teacher machine learning model, and instead includes a different feature extractor having fewer parameters than the pre-trained feature extractor.