Single-stage model training for neural architecture search
Methods, systems, and apparatus, including computer programs encoded on computer storage media, for selecting a neural network to perform a particular machine learning task while satisfying a set of constraints.
1 . A method performed by one or more computers, the method comprising:
receiving training data for a particular machine learning task;
generating trained values for a shared set of parameters by training a plurality of neural networks jointly on the training data, wherein each of the plurality of neural networks has parameters that are a subset of the shared set, wherein each of the plurality of neural networks has a respective architecture selected from a search space of different architectures, and wherein the search space is defined by a respective set of possible values for each of a plurality of architectural dimensions, and wherein the joint training comprises performing the following at each of a plurality of training steps:
obtaining training data for the training step;
sampling a plurality of neural networks of different sizes from the plurality of neural networks being trained, wherein each sampled neural network has a different subset of the shared set of parameters;
training the sampled neural networks on the training data for the training step to determine, for each sampled neural network, gradients with respect the respective subset of the shared set of parameters that are included in the sampled neural network; and
aggregating the gradients determined for the set of neural networks to determine an update to the shared set of parameters, comprising, for each parameter that is shared between two or more of the sampled neural networks, aggregating the gradients with respect to the parameter for the two or more sampled neural networks;
determining a respective performance benchmark on the machine learning task for each of a subset of neural networks, each neural network in the subset having a different architecture selected from a proper subset of the search space of different architectures;
receiving data specifying one or more constraints on computational resource consumption for performing inference on one or more inputs for the particular machine learning task on a particular set of one or more target hardware devices; and
selecting a neural network to be deployed for performing the machine learning task, comprising:
identifying, from the neural networks in the subset that satisfy the one or more constraints, an initial neural network based on the respective performance benchmarks;
generating a plurality of modified neural networks by perturbing the architecture of the identified initial neural network;
determining a respective performance benchmark on the particular machine learning task for each of the modified neural networks, and
selecting, from the identified neural network and the plurality of modified neural networks generated by perturbing the architecture of the identified initial neural network, a neural network based on the respective performance benchmarks; and
providing the selected neural network for performing inference on the particular set of one or more target hardware devices.
2 . The method of claim 1 , wherein the one or more constraints comprise one or more constraints on one or more of: floating point operations per second performed by a device when performing the task using the neural network, a memory footprint of the neural network when deployed on a device for performing the task, or a runtime latency of a device when performing an inference for an input or a batch of inputs using the neural network.
3 . The method of claim 1 , wherein the one or more constraints specify one or more of a maximum number of floating point operations per second or an acceptable range of memory consumed by a device when performing the task using the neural network, or a maximum latency or an acceptable latency range of a device when performing an inference for an input or a batch of inputs using the neural network.
4 . The method of claim 1 , wherein the constraint is a hardware-specific constraint on computational resource consumption for a particular device.
5 . The method of claim 1 , wherein providing the selected neural network for performing inference on the particular set of one or more target hardware devices comprises deploying the selected neural network with values of the parameters of the selected neural network that are the same as the trained values determined by training the plurality of neural networks jointly.
6 . The method of claim 1 , wherein the selected neural network has an architecture defined by respective first values for each of the plurality of architectural dimensions and wherein generating a plurality of modified neural networks comprises:
randomly adjusting one or more of the respective first values to generate a modified architecture for a modified neural network.
7 . The method of claim 1 wherein the sampled networks comprise a largest possible neural network in the search space, a smallest possible neural network in the search space, and N randomly sampled neural networks from the search space.
8 . The method of claim 7 , wherein training the set of neural networks comprises:
training the largest possible neural network using target outputs from the training data; and
training each neural network in the set other than the largest possible neural network using, as target outputs, outputs generated by the largest possible neural network during the training.
9 . The method of claim 7 , wherein training the set of neural networks comprises:
regularizing the largest possible neural network without regularizing any of the other neural networks in the set.
10 . The method of claim 7 wherein the largest possible neural network has all of the parameters in the shared set of parameters.
11 . The method of claim 7 wherein the smallest possible neural network has a smallest number of parameters of any neural network having an architecture in the search space.
12 . The method of claim 1 , wherein the respective architecture of each of the plurality of neural networks comprises one or more residual blocks, and wherein training the plurality of neural networks comprises:
initializing an output of each residual block to an all-zeros tensor.
13 . The method of claim 1 , wherein training the plurality of neural networks comprises:
training using a learning rate schedule that is exponentially decaying with constant ending.
14 . The method of claim 1 , wherein providing the selected neural network for performing inference on the particular set of one or more target hardware devices comprises deploying the selected neural network on the particular set of one or more target hardware devices without performing any additional training.
15 . 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:
receiving training data for a particular machine learning task;
generating trained values for a shared set of parameters by training a plurality of neural networks jointly on the training data, wherein each of the plurality of neural networks has parameters that are a subset of the shared set, wherein each of the plurality of neural networks has a respective architecture selected from a search space of different architectures, and wherein the search space is defined by a respective set of possible values for each of a plurality of architectural dimensions, and wherein the joint training comprises performing the following at each of a plurality of training steps:
obtaining training data for the training step;
sampling a plurality of neural networks of different sizes from the plurality of neural networks being trained, wherein each sampled neural network has a different subset of the shared set of parameters;
training the sampled neural networks on the training data for the training step to determine, for each sampled neural network, gradients with respect the respective subset of the shared set of parameters that are included in the sampled neural network; and
aggregating the gradients determined for the set of neural networks to determine an update to the shared set of parameters, comprising, for each parameter that is shared between two or more of the sampled neural networks, aggregating the gradients with respect to the parameter for the two or more sampled neural networks;
determining a respective performance benchmark on the machine learning task for each of a subset of neural networks, each neural network in the subset having a different architecture selected from a proper subset of the search space of different architectures;
receiving data specifying one or more constraints on computational resource consumption for performing inference on one or more inputs for the particular machine learning task on a particular set of one or more target hardware devices; and
selecting a neural network to be deployed for performing the machine learning task, comprising:
identifying, from the neural networks in the subset that satisfy the one or more constraints, an initial neural network based on the respective performance benchmarks;
generating a plurality of modified neural networks by perturbing the architecture of the identified initial neural network;
determining a respective performance benchmark on the particular machine learning task for each of the modified neural networks, and
selecting, from the identified neural network and the plurality of modified neural networks generated by perturbing the architecture of the identified initial neural network, a neural network based on the respective performance benchmarks; and
providing the selected neural network for performing inference on the particular set of one or more target hardware devices.
16 . A system comprising one or more computers and one or more storage devices storing instructions that when executed by one or more computers are operable to cause the one or more computers to perform operations comprising:
receiving training data for a particular machine learning task;
generating trained values for a shared set of parameters by training a plurality of neural networks jointly on the training data, wherein each of the plurality of neural networks has parameters that are a subset of the shared set, wherein each of the plurality of neural networks has a respective architecture selected from a search space of different architectures, and wherein the search space is defined by a respective set of possible values for each of a plurality of architectural dimensions, and wherein the joint training comprises performing the following at each of a plurality of training steps:
obtaining training data for the training step;
sampling a plurality of neural networks of different sizes from the plurality of neural networks being trained, wherein each sampled neural network has a different subset of the shared set of parameters;
training the sampled neural networks on the training data for the training step to determine, for each sampled neural network, gradients with respect the respective subset of the shared set of parameters that are included in the sampled neural network; and
aggregating the gradients determined for the set of neural networks to determine an update to the shared set of parameters, comprising, for each parameter that is shared between two or more of the sampled neural networks, aggregating the gradients with respect to the parameter for the two or more sampled neural networks;
determining a respective performance benchmark on the machine learning task for each of a subset of neural networks, each neural network in the subset having a different architecture selected from a proper subset of the search space of different architectures;
receiving data specifying one or more constraints on computational resource consumption for performing inference on one or more inputs for the particular machine learning task on a particular set of one or more target hardware devices; and
selecting a neural network to be deployed for performing the machine learning task, comprising:
identifying, from the neural networks in the subset that satisfy the one or more constraints, an initial neural network based on the respective performance benchmarks;
generating a plurality of modified neural networks by perturbing the architecture of the identified initial neural network;
determining a respective performance benchmark on the particular machine learning task for each of the modified neural networks, and
selecting, from the identified neural network and the plurality of modified neural networks generated by perturbing the architecture of the identified initial neural network, a neural network based on the respective performance benchmarks; and
providing the selected neural network for performing inference on the particular set of one or more target hardware devices.
17 . The system of claim 16 , wherein the one or more constraints comprise one or more constraints on one or more of: floating point operations per second performed by a device when performing the task using the neural network, a memory footprint of the neural network when deployed on a device for performing the task, or a runtime latency of a device when performing an inference for an input or a batch of inputs using the neural network.
18 . The system of claim 16 , wherein the one or more constraints specify one or more of a maximum number of floating point operations per second or an acceptable range of memory consumed by a device when performing the task using the neural network, or a maximum latency or an acceptable latency range of a device when performing an inference for an input or a batch of inputs using the neural network.
19 . The system of claim 16 , wherein the constraint is a hardware-specific constraint on computational resource consumption for a particular device.
20 . The system of claim 16 , wherein providing the selected neural network for performing inference on the particular set of one or more target hardware devices comprises deploying the selected neural network with values of the parameters of the selected neural network that are the same as the trained values determined by training the plurality of neural networks jointly.