Neural architecture search for dense image prediction tasks

A system comprising one or more computers and one or more storage devices storing instructions that when executed by the one or more computers cause the one or more computers to perform operations comprising:

• obtaining training data for a dense image prediction task, wherein the dense image prediction task is a task that requires processing an image to generate an output that includes a respective prediction for each of a plurality of pixels in the image, and wherein the training data comprises a plurality of training inputs;
• and
• determining an architecture for a neural network configured to perform the dense image prediction task, comprising:
• searching a space of candidate architectures to identify one or more best performing architectures using the training data, wherein the candidate architectures in the space of candidate architectures all have (i) a same first neural network backbone that is configured to receive an input image and to process the input image to generate a plurality of feature maps and (ii) a respective dense prediction cell configured to process the plurality of feature maps and to generate an output for the dense image prediction task, wherein each candidate architecture in the space has a different respective dense prediction cell from each other candidate architecture in the space, wherein the searching comprises:
• obtaining data specifying a pre-trained first neural network backbone, the data specifying the pre-trained first neural network backbone comprising trained values of the parameters of the first neural network backbone;
• processing each of the plurality of training inputs using the pre-trained first neural network backbone and in accordance with the trained values of the parameters of the first neural network backbone to generate a respective plurality of feature maps for each of the plurality of training inputs;
• storing the respective plurality of feature maps for each of the plurality of training inputs in one or more memories;
• and
• repeatedly performing the following operations:
• selecting one or more candidate architectures from the space of candidate architectures, each selected candidate architecture having the same first neural network backbone and a respective dense prediction cell;
• training each of the selected one or more candidate architectures on at least a portion of the training data to update parameters of the respective dense prediction cell in the candidate architecture while holding the trained values of the parameters of the first neural network backbone fixed, the training comprising, for each training input in the portion of the training data:
• accessing the respective plurality of feature maps generated for the training input by the pre-trained dense prediction cell from the one or more memories;
• and
• providing the plurality of feature maps that were accessed from memory as input to the respective dense prediction cell instead of processing the training input using the first neural network backbone;
• for each of the trained candidate architectures, evaluating the performance of the trained candidate architecture on the dense image prediction task;
• and
• determining the architecture for the neural network based on the best performing candidate architectures.

One or more computer-readable storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations comprising:

• obtaining training data for a dense image prediction task, wherein the dense image prediction task is a task that requires processing an image to generate an output that includes a respective prediction for each of a plurality of pixels in the image, and wherein the training data comprises a plurality of training inputs;
• and
• determining an architecture for a neural network configured to perform the dense image prediction task, comprising:
• searching a space of candidate architectures to identify one or more best performing architectures using the training data, wherein the candidate architectures in the space of candidate architectures all have (i) a same first neural network backbone that is configured to receive an input image and to process the input image to generate a plurality of feature maps and (ii) a respective dense prediction cell configured to process the plurality of feature maps and to generate an output for the dense image prediction task, wherein each candidate architecture in the space has a different respective dense prediction cell from each other candidate architecture in the space, wherein the searching comprises:
• obtaining data specifying a pre-trained first neural network backbone, the data specifying the pre-trained first neural network backbone comprising trained values of the parameters of the first neural network backbone;
• processing each of the plurality of training inputs using the pre-trained first neural network backbone and in accordance with the trained values of the parameters of the first neural network backbone to generate a respective plurality of feature maps for each of the plurality of training inputs;
• storing the respective plurality of feature maps for each of the plurality of training inputs in one or more memories;
• and
• repeatedly performing the following operations:
• selecting one or more candidate architectures from the space of candidate architectures, each selected candidate architecture having the same first neural network backbone and a respective dense prediction cell;
• training each of the selected one or more candidate architectures on at least a portion of the training data to update parameters of the respective dense prediction cell in the candidate architecture while holding the trained values of the parameters of the first neural network backbone fixed, the training comprising, for each training input in the portion of the training data:
• accessing the respective plurality of feature maps generated for the training input by the pre-trained dense prediction cell from the one or more memories;
• and
• providing the plurality of feature maps that were accessed from memory as input to the respective dense prediction cell instead of processing the training input using the first neural network backbone;
• for each of the trained candidate architectures, evaluating the performance of the trained candidate architecture on the dense image prediction task;
• and
• determining the architecture for the neural network based on the best performing candidate architectures.