Neural architecture search for dense image prediction tasks

Abstract

Methods, systems, and apparatus, including computer programs encoded on computer storage media, for determining neural network architectures. One of the methods includes obtaining training data for a dense image prediction task; 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 each candidate architecture in the space of candidate architectures comprises (i) the 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 different dense prediction cell configured to process the plurality of feature maps and to generate an output for the dense image prediction task; and determining the architecture for the neural network based on the best performing candidate architectures.

1 . (canceled) 
     
     
         2 . A method 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;   identifying, using the training data, one or more best performing architectures for a neural network that is configured to perform dense image prediction task, wherein the identifying 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 and the pre-trained neural network backbone being configured to receive an input image and to process the input image to generate a plurality of feature maps; 
 repeatedly performing the following operations:
 selecting one or more candidate architectures from a space of candidate architectures, each selected candidate architecture having the pre-trained first neural network backbone and a different dense prediction cell, each dense prediction cell configured to process the plurality of feature maps and to generate an output for the dense image prediction task; 
 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; and 
 for each of the trained candidate architectures, evaluating the performance of the trained candidate architecture on the dense image prediction task; and 
 
   determining a final architecture for the neural network based on the identified best performing architectures.   
     
     
         3 . The method of  claim 2 , wherein determining a final architecture for the neural network based on the identified best performing architectures comprises:
 selecting one of the best performing architectures; and   generating a new architecture by replacing the first neural network backbone with a second, different neural network backbone that has more parameters than the first neural network backbone.   
     
     
         4 . The method of  claim 2 , wherein training each of the selected one or more candidate architectures comprises training each selected candidate architecture for a fixed number of iterations instead of to convergence. 
     
     
         5 . The method of  claim 2 , further comprising:
 pre-training the first neural network backbone to determine the trained values of the parameters of the first neural network backbone.   
     
     
         6 . The method of  claim 2 , wherein selecting one or more candidate architectures comprises:
 selecting the one or more candidate architectures using a random search strategy.   
     
     
         7 . The method of  claim 2 , wherein determining the final architecture for the neural network based on the best performing architectures comprises:
 for each of the one or more best performing architectures, generating a final architecture and further training the final architecture to convergence on the dense prediction task using the training data; and   selecting the best performing final architecture as the architecture for the neural network.   
     
     
         8 . The method of  claim 7 , wherein generating the final architecture comprises replacing the first neural network backbone with a second, different neural network backbone having more parameters than the first neural network backbone. 
     
     
         9 . The method of  claim 2 , wherein each respective dense prediction cell includes B branches, wherein B is a fixed integer greater than one, and wherein each of the B branches maps an input tensor to an output tensor by applying an operation to the input tensor. 
     
     
         10 . The method of  claim 9 , wherein the dense prediction cell combines the output tensors generated by the B branches to generate a combined output tensor. 
     
     
         11 . The method of  claim 10 , wherein the dense prediction cell concatenates the output tensors. 
     
     
         12 . The method of  claim 10 , wherein the dense prediction cell processes the combined output tensor through one or more output layers to generate the output for the dense prediction task. 
     
     
         13 . The method of  claim 9 , wherein each candidate architecture specifies (i) an input tensor from a set of input tensors to be provided as input to each of the B branches in the respective dense prediction cell in the candidate architecture, and (ii) an operation from a set of operations to be performed by each of the B branches in the respective dense prediction cell in the candidate architecture. 
     
     
         14 . The method of  claim 13 , wherein the set of operations includes a convolution with a 1×1 kernel. 
     
     
         15 . The method of  claim 13 , wherein the set of operations includes one or more atrous separable convolution operations, each having a respective sampling rate. 
     
     
         16 . The method of  claim 13 , wherein the set of operations includes one or more spatial pyramid pooling operations, each having a respective grid size. 
     
     
         17 . The method of  claim 9 , wherein the B branches are ordered and wherein, for each of the B branches, the set of input tensors includes (i) the feature maps generated by the first network backbone and (ii) output tensors generated by any branches before the branch in the order of the B branches. 
     
     
         18 . One or more non-transitory 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;   identifying, using the training data, one or more best performing architectures for a neural network that is configured to perform dense image prediction task, wherein the identifying 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 and the pre-trained neural network backbone being configured to receive an input image and to process the input image to generate a plurality of feature maps; 
 repeatedly performing the following operations:
 selecting one or more candidate architectures from a space of candidate architectures, each selected candidate architecture having the pre-trained first neural network backbone and a different dense prediction cell, each dense prediction cell configured to process the plurality of feature maps and to generate an output for the dense image prediction task; 
 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; and 
 for each of the trained candidate architectures, evaluating the performance of the trained candidate architecture on the dense image prediction task; and 
 
   determining a final architecture for the neural network based on the identified best performing architectures.   
     
     
         19 . The computer-readable storage media of  claim 18 , wherein determining a final architecture for the neural network based on the identified best performing architectures comprises:
 selecting one of the best performing architectures; and   generating a new architecture by replacing the first neural network backbone with a second, different neural network backbone that has more parameters than the first neural network backbone.   
     
     
         20 . 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;   identifying, using the training data, one or more best performing architectures for a neural network that is configured to perform dense image prediction task, wherein the identifying 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 and the pre-trained neural network backbone being configured to receive an input image and to process the input image to generate a plurality of feature maps; 
 repeatedly performing the following operations:
 selecting one or more candidate architectures from a space of candidate architectures, each selected candidate architecture having the pre-trained first neural network backbone and a different dense prediction cell, each dense prediction cell configured to process the plurality of feature maps and to generate an output for the dense image prediction task; 
 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; and 
 for each of the trained candidate architectures, evaluating the performance of the trained candidate architecture on the dense image prediction task; and 
 
   determining a final architecture for the neural network based on the identified best performing architectures.   
     
     
         21 . The system of  claim 20 , wherein determining a final architecture for the neural network based on the identified best performing architectures comprises:
 selecting one of the best performing architectures; and   generating a new architecture by replacing the first neural network backbone with a second, different neural network backbone that has more parameters than the first neural network backbone.