Desparsified convolution for sparse activations
Certain aspects of the present disclosure provide techniques for desparsified convolution. An activation tensor is received, and a convolution output is generated for the activation tensor, comprising: selecting a subset of weight elements, corresponding to a set of non-zero elements in the activation tensor, from a weight tensor, and multiplying the set of non-zero elements and the set of weight elements.
1 . A processor-implemented method, comprising:
receiving an input to a machine learning model;
generating a first convolution output based on the input, comprising:
generating an activation tensor based on the input;
selecting a first subset of weight elements, corresponding to a first set of non-zero elements in the activation tensor, from a weight tensor; and
multiplying the first set of non-zero elements and the first subset of weight elements; and
generating an output of the machine learning model based on the first convolution output.
2 . The method of claim 1 , further comprising identifying a set of respective indices, wherein:
each respective index of the set of respective indices corresponds to a respective non-zero element of the first set of non-zero elements, and
selecting the first subset of weight elements comprises retrieving the first subset of weight elements based on the set of respective indices.
3 . The processor-implemented method of claim 2 , wherein retrieving the first subset of weight elements based on the set of respective indices comprises:
determining a set of memory offsets based on the set of respective indices; and
retrieving the first subset of weight elements from a memory using the set of memory offsets.
4 . The processor-implemented method of claim 1 , wherein:
the first set of non-zero elements is identified within a first one-dimensional vector extracted from the activation tensor, and
the first one-dimensional vector corresponds to a single spatial element across input channels in the activation tensor.
5 . The processor-implemented method of claim 4 , wherein:
the first convolution output corresponds to a first application of a kernel having a height of M and a width of N, and
generating the first convolution output comprises:
identifying M*N sets of non-zero elements in the activation tensor;
selecting M*N subsets of weight elements corresponding to the M*N sets of non-zero elements; and
generating a set of intermediate tensors by multiplying each of the M*N subsets of weight elements with a corresponding set of non-zero elements from the M*N sets of non-zero elements.
6 . The processor-implemented method of claim 5 , wherein generating the first convolution output further comprises summing the set of intermediate tensors to yield the first convolution output.
7 . The processor-implemented method of claim 5 , further comprising:
generating a second convolution output corresponding to a second application of the kernel; and
generating an overall convolution output for the activation tensor by stacking the first and second convolution outputs.
8 . The processor-implemented method of claim 7 , wherein:
H*W convolution outputs are generated for the activation tensor,
H is a height of the activation tensor, and
W is a width of the activation tensor.
9 . The processor-implemented method of claim 1 , wherein the first convolution output is generated as part of applying a convolution kernel comprising at least one of:
a width dimension greater than one, or
a height dimension greater than one.
10 . The processor-implemented method of claim 1 , wherein the first convolution output is generated as part of applying a convolution kernel comprising:
a width dimension equal to one, and
a height dimension equal to one.
11 . The processor-implemented method of claim 1 , wherein the output is generated as part of a training process for a neural network.
12 . The processor-implemented method of claim 1 , wherein the output is generated as part of generating an inference using a trained neural network.
13 . A system, comprising:
memory comprising computer-executable instructions; and
one or more processors configured to execute the computer-executable instructions and cause the system to perform an operation comprising:
receiving an input to a machine learning model;
generating a first convolution output based on the input for the activation tensor, comprising:
generating an activation tensor based on the input;
selecting a first subset of weight elements, corresponding to a first set of non-zero elements in the activation tensor, from a weight tensor; and
multiplying the first set of non-zero elements and the first subset of weight elements; and
generating an output of the machine learning model based on the first convolution output.
14 . The system of claim 13 , the operation further comprising identifying a set of respective indices, wherein:
each respective index of the set of respective indices corresponds to a respective non-zero element of the first set of non-zero elements, and
selecting the first subset of weight elements comprises retrieving the first subset of weight elements based on the set of respective indices.
15 . The system of claim 14 , wherein retrieving the first subset of weight elements based on the set of respective indices comprises:
determining a set of memory offsets based on the set of respective indices; and
retrieving the first subset of weight elements from a memory using the set of memory offsets.
16 . The system of claim 15 , wherein the memory is a tightly coupled memory (TCM) having a dedicated connection to a processing core used to generate the first convolution output.
17 . The system of claim 13 , wherein:
the first set of non-zero elements are identified within a first one-dimensional vector extracted from the activation tensor, and
the first one-dimensional vector corresponds to a single spatial element across input channels in the activation tensor.
18 . The system of claim 17 , wherein:
the first convolution output corresponds to a first application of a kernel having a height of M and a width of N, and
generating the first convolution output comprises:
identifying M*N sets of non-zero elements in the activation tensor;
selecting M*N subsets of weight elements corresponding to the M*N sets of non-zero elements; and
generating a set of intermediate tensors by multiplying each of the M*N subsets of weight elements with a corresponding set of non-zero elements from the M*N sets of non-zero elements.
19 . The system of claim 18 , wherein generating the first convolution output further comprises summing the set of intermediate tensors to yield the first convolution output.
20 . The system of claim 18 , the operation further comprising:
generating a second convolution output corresponding to a second application of the kernel; and
generating an overall convolution output for the activation tensor by stacking the first and second convolution outputs.
21 . The system of claim 13 , wherein the first convolution output is generated as part of applying a convolution kernel comprising at least one of:
a width dimension greater than one, or
a height dimension greater than one.
22 . The system of claim 13 , wherein the first convolution output is generated as part of applying a convolution kernel comprising:
a width dimension equal to one, and
a height dimension equal to one.
23 . One or more non-transitory computer-readable media comprising instructions that, when executed by one or more processors of a processing system, cause the processing system to perform an operation comprising:
receiving an input to a machine learning model;
generating a first convolution output based on the input, comprising:
generating an activation tensor based on the input;
selecting a first subset of weight elements, corresponding to a first set of non-zero elements in the activation tensor, from a weight tensor; and
multiplying the first set of non-zero elements and the first subset of weight elements; and
generating an output of the machine learning model based on the first convolution output.
24 . The non-transitory computer-readable media of claim 23 , the operation further comprising identifying a set of respective indices, wherein:
each respective index of the set of respective indices corresponds to a respective non-zero element of the first set of non-zero elements, and
selecting the first subset of weight elements comprises retrieving the first subset of weight elements based on the set of respective indices.
25 . The non-transitory computer-readable media of claim 24 , wherein retrieving the first subset of weight elements based on the set of respective indices comprises:
determining a set of memory offsets based on the set of respective indices; and
retrieving the first subset of weight elements from a memory using the set of memory offsets.
26 . The non-transitory computer-readable media of claim 23 , wherein:
the first set of non-zero elements are identified within a first one-dimensional vector extracted from the activation tensor, and
the first one-dimensional vector corresponds to a single spatial element across input channels in the activation tensor.
27 . The non-transitory computer-readable media of claim 26 , wherein:
the first convolution output corresponds to a first application of a kernel having a height of M and a width of N, and
generating the first convolution output comprises:
identifying M*N sets of non-zero elements in the activation tensor;
selecting M*N subsets of weight elements corresponding to the M*N sets of non-zero elements; and
generating a set of intermediate tensors by multiplying each of the M*N subsets of weight elements with a corresponding set of non-zero elements from the M*N sets of non-zero elements.
28 . The non-transitory computer-readable media of claim 27 , wherein generating the first convolution output further comprises summing the set of intermediate tensors to yield the first convolution output.
29 . The non-transitory computer-readable media of claim 28 , wherein generating the first convolution output further comprises summing the set of intermediate tensors to yield the first convolution output.
30 . A system, comprising:
means for receiving an input to a machine learning model;
means for generating a convolution output based on the input, comprising:
means for generating an activation tensor based on the input;
means for selecting a subset of weight elements, corresponding to a set of non-zero elements in the activation tensor, from a weight tensor; and
means for multiplying the set of non-zero elements and the subset of weight elements; and
means for generating an output of the machine learning model based on the convolution output.