Training optimization for low memory footprint
A method is described. The method includes profiling a model for a learning network having a plurality of layers and associated memory. The layers include weight layers and activation layers. The plurality of weight layers including weights. The method also includes determining, based on the profiling, a training technique for the model on the learning network. The determination of the training technique includes optimizing a latency for at least one training iteration for a capacity of the associated memory. The training iteration(s) include at least one update of the weights for the weight layers.
1 . A method, comprising:
profiling a model for a learning network having a plurality of layers and associated memory, the plurality of layers including a plurality of weight layers and a plurality of activation layers, the plurality of weight layers including a plurality of weights, the plurality of weight layers corresponding to at least one compute engine, each of the at least one compute engine including a compute-in-memory (CIM) hardware module, the CIM hardware module being configured to store a least a portion of the plurality of weights and to perform a hardware vector-matrix multiplication (VMM) for the portion of the plurality of weights stored by the CIM hardware module, the associated memory being coupled with the at least one compute engine, the profiling including at least one of:
performing, on the learning network, at least one training iteration for the model using a particular training technique, the performing the at least one training iteration including propagating an activation through the plurality of weight layers and through the plurality of activation layers to provide an output and determining a weight update based upon the output, a target output, and the particular training technique, the propagating the activation through the plurality of weight layers including performing the hardware VMM by the CIM hardware module for each of the at least one compute engine; or
determining a plurality of descriptors for the model executed using the at least one compute engine for selection, based on the plurality of descriptors, of at least a portion of a training technique from a library including a plurality of training techniques, the plurality of training techniques in the library being determined by executing the plurality of training techniques on at least an other compute engine that is analogous to the at least one compute engine based on the plurality of descriptors, the executing the plurality of training techniques on the at least the other compute engine including, for a selected training technique of the plurality of training techniques, propagating an other activation through an other plurality of weight layers and through an other plurality of activation layers to provide an other output and determining an other weight update based upon the other output, an other target output, and the selected training technique, the other plurality of weight layers including an other plurality of weights stored in a hardware compute engine of the at least the other compute engine; and
determining, based on the profiling, the training technique for the model on the learning network, the determining including optimizing a latency for the at least one training iteration for a capacity of the associated memory; and
wherein the training technique is used to update the plurality of weights stored in the CIM hardware module during training of the learning network.
2 . The method of claim 1 , wherein the particular training technique includes at least one of backpropagation, activation recomputation, partial activation recomputation, or feedback alignment.
3 . The method of claim 2 , wherein the performing the training iteration further includes:
performing an inference using a training activation input;
determining a loss function based on the inference;
using the loss function and the particular training technique to determine a plurality of updates for the plurality of weights; and
updating the plurality of weights.
4 . The method of claim 3 , wherein the using the loss function to determine the plurality of updates further includes:
backpropagating the loss function through the plurality of layers.
5 . The method of claim 1 , wherein the capacity of the associated memory is less than a memory required for the training iteration, performing the training iteration further includes:
loading, from a remote memory, the activation for a weight layer of the plurality of weight layers;
storing an output of the weight layer in the remote memory; and
repeating the loading and storing for each subsequent weight layer of the plurality of weight layers.
6 . The method of claim 1 , wherein the profiling further includes:
the determining the plurality of descriptors for the model; and
based on the plurality of descriptors, selecting the at least a portion of the training technique from the library including the plurality of training techniques.
7 . The method of claim 6 , wherein the plurality of descriptors includes at least one of a batch size, a channel number, a width of an input activation, a height of the input activation, a stride, a padding size, a kernel size, a pooling size, or a dilation.
8 . The method of claim 1 , further comprising:
performing the training of the learning network, the training including propagating a training activation through the plurality of weight layers and through the plurality of activation layers to provide a training output, determining a training weight update based upon the training output, a training target output, and the training technique, and updating the plurality of weights in the at least one compute engine using the training weight update
wherein the training technique is used to update the plurality of weights stored in the CIM hardware module during training of the learning network.
9 . The method of claim 1 , wherein the training of the learning network does not exceed the capacity of the memory associated with the learning network.
10 . A system, comprising:
a memory; and
a processor coupled to the memory and configured to:
profile a model for a learning network having a plurality of layers and associated memory, the plurality of layers including a plurality of weight layers and a plurality of activation layers, the plurality of weight layers including a plurality of weights, the plurality of weight layers corresponding to at least one compute engine, each of the at least one compute engine including a compute-in-memory (CIM) hardware module, the CIM hardware module being configured to store a least a portion of the plurality of weights and to perform a hardware vector-matrix multiplication (VMM) of an input activation and the portion of the plurality of weights stored by the CIM hardware module, the associated memory being coupled with the at least one compute engine, wherein to profile the processor is further configured to at least one of
manage performing, on the learning network, at least one training iteration for the model using a particular training technique using the at least one compute engine, the performing the at least one training iteration including propagating an activation through the plurality of weight layers and through the plurality of activation layers to provide an output and determining a weight update based upon the output, a target output, and the particular training technique, the propagating the activation through the plurality of weight layers including performing the hardware VMM by the CIM hardware module for each of the at least one compute engine, or
or determine a plurality of descriptors for the model executed using the at least one compute engine based on the plurality of descriptors for selection, based on the plurality of descriptors, of at least a portion of a training technique from a library including a plurality of training techniques, the plurality of training techniques in the library being determined by executing the plurality of training techniques on at least an other compute engine that is analogous to the at least one compute engine based on the plurality of descriptors, the executing the plurality of training techniques on the at least the other compute engine including, for a selected training technique of the plurality of training techniques, propagating an other activation through an other plurality of weight layers and through an other plurality of activation layers to provide an other output and determining an other weight update based upon the other output, an other target output, and the selected training technique, the other plurality of weight layers including an other plurality of weights stored in a hardware compute engine of the at least the other compute engine; and
determine, based on the profiling, the training technique for the model on the learning network, to determine, the processor being further configured to optimize a latency for the at least one training iteration for a capacity of the associated memory;
wherein the training technique is used to update the plurality of weights stored in the CIM hardware module during training of the learning network.
11 . The system of claim 9 , wherein the particular training technique includes at least one of backpropagation, activation recomputation, partial activation recomputation, or feedback alignment.
12 . The system of claim 11 , wherein the performing the at least one training iteration further includes the learning network:
performing an inference using a training activation input;
determining a loss function based on the inference;
using the loss function and the particular training technique to determine the weight update; and
updating the plurality of weights.
13 . The system of claim 12 , wherein the using the loss function to determine the plurality of updates further includes:
backpropagating the loss function through the plurality of layers.
14 . The system of claim 9 , wherein the capacity of the associated memory is less than a memory required for the training iteration, performing the training iteration further includes:
loading, from a remote memory, an activation for a weight layer of the plurality of weight layers;
storing an output of the weight layer in the remote memory; and
repeating the loading and storing for each subsequent weight layer of the plurality of weight layers.
15 . The system of claim 10 , wherein the processor being configured to profile further includes the processor being configured to:
determine a plurality of descriptors for the model; and
based on the plurality of descriptors, select at least a portion of the training technique from a library including a plurality of training techniques.
16 . The system of claim 15 , wherein the plurality of descriptors includes at least one of a batch size, a channel number, a width of an input activation, a height of the input activation, a stride, a padding size, a kernel size, a pooling size, or a dilation.
17 . The system of claim 10 , wherein the learning network includes at least one tile, each of the at least one tile including:
the at least one compute engine, each of the at least one compute engine including the compute-in-memory (CIM) hardware module, the portion of the plurality of weights corresponding to a matrix and the hardware VMM being for the matrix; and
a general-purpose (GP) processor coupled with the plurality of compute engines and configured to provide control instructions and data to the plurality of compute engines, the general-purpose processor corresponding to the plurality of activation layers.
18 . The system of claim 17 , wherein the GP processor is coupled with the plurality of compute engines via a streaming port and a control port, the control port configured to provide the plurality of compute engines with the control instructions and the streaming port being configured to exchange data between the GP processor and the plurality of compute engines.
19 . The system of claim 17 , wherein the CIM hardware module includes a plurality of storage cells for storing the plurality of weights.
20 . A computer program product embodied in a non-transitory computer readable medium and comprising computer instructions for:
profiling a model for a learning network having a plurality of layers and associated memory, the plurality of layers including a plurality of weight layers and a plurality of activation layers, the plurality of weight layers including a plurality of weights, the plurality of weight layers corresponding to at least one compute engine, each of the at least one compute engine including a compute-in-memory (CIM) hardware module, the CIM hardware module being configured to store a least a portion of the plurality of weights and to perform a hardware vector-matrix multiplication (VMM) of an input activation and the portion of the plurality of weights stored by the CIM hardware module, the associated memory being coupled with the at least one compute engine, the computer instructions for profiling including computer instructions for at least one of:
performing, on the learning network, at least one training iteration for the model using a particular training technique, the performing the at least one training iteration including propagating an activation through the plurality of weight layers and through the plurality of activation layers to provide an output and determining a weight update based upon the output, a target output, and the particular training technique, the propagating the activation through the plurality of weight layers including performing the hardware VMM by the CIM hardware module for each of the at least one compute engine; or
determining a plurality of descriptors for the model executed using the at least one compute engine for selection, based on the plurality of descriptors, of at least a portion of a training technique from a library including a plurality of training techniques, the plurality of training techniques in the library being determined by executing the plurality of training techniques on at least an other compute engine that is analogous to the at least one compute engine
based on the plurality of descriptors, the executing the plurality of training techniques on the at least the other compute engine including, for a selected training technique of the plurality of training techniques, propagating an other activation through an other plurality of weight layers and through an other plurality of activation layers to provide an other output and determining an other weight update based upon the other output, an other target output, and the selected training technique, the other plurality of weight layers including an other plurality of weights stored in a hardware compute engine of the at least the other compute engine; and
determining, based on the profiling, the training technique for the model on the learning network, the determining including optimizing a latency for the at least one training iteration for a capacity of the associated memory;
wherein the training technique is used to update the plurality of weights stored in the CIM hardware module during training of the learning network.