Method and system for balanced-weight sparse convolution processing

1. A computer-implemented method, comprising:

obtaining an input tensor and a plurality of filters at a layer within a neural network;

segmenting the input tensor into a plurality of sub-tensors;

for each of the plurality of filters, pruning the filter, wherein the pruning comprises:

dividing a channel dimension of the filter into a plurality of channel groups, each channel group comprising a same number of two or more continuous channels of the filter, and

pruning the filter so that (1) each channel group within the filter comprises a same number of one or more non-zero weights, and (2) all non-zero weights in the filter are evenly distributed in the plurality of channel groups within the filter, wherein the pruning comprises setting one or more non-zero weights to zeros;

storing the non-zero weights in each of the plurality of channel groups of each filter as index-value pairs;

segmenting, according to the plurality of pruned channel groups, each of the plurality of filters into a plurality of sub-filters, wherein each of the plurality of sub-filters comprises the same number of non-zero weights and the corresponding index-value pair consumes a same number of bits; and

assigning the plurality of sub-tensors and the plurality of sub-filters to a plurality of processors for parallel convolution processing.

2. The method of claim 1 , wherein the method further comprises performing the parallel convolution processing by:

iterating, by each of the plurality of processors, each non-zero weight of a sub-filter assigned to the processor; and

identifying, by the processor, a corresponding input value in a sub-tensor assigned to the processor to perform a multiply-and-accumulate (MAC) operation.

3. The method of claim 1 , wherein the method further comprises performing the parallel convolution processing by:

for each of the non-zero weights stored as an index-value pair, identifying a corresponding input value in an assigned sub-tensor at a location identified by the channel-dimension index and the width-dimension index of the index-value pair representing each non-zero weight.

4. The method of claim 1 , further comprising:

rotating the plurality of sub-filters among the plurality of processors.

5. The method of claim 1 , wherein segmenting each of the plurality of filters into the plurality of sub-filters according to the plurality of pruned channel groups comprises:

segmenting the filter into a plurality of chunks according to the plurality of pruned channel groups; and

segmenting each of the plurality of chunks into a plurality of horizontal planes.

6. The method of claim 5 , further comprising:

pruning the plurality of horizontal planes so that each of the plurality of horizontal planes comprises the same number of non-zero weights.

7. The method of claim 1 , wherein after the parallel convolution processing, the plurality of processors generate a plurality of partial sums, and

the method further comprises:

accumulating the plurality of partial sums to obtain an output tensor; and

feeding the output tensor as an input tensor for a next layer of the neural network.

8. A system, comprising:

a plurality of processors; and

one or more non-transitory computer-readable memories coupled to the plurality of processors and configured with instructions executable by the plurality of processors to cause the system to perform operations comprising:

obtaining an input tensor and a plurality of filters at a layer within a neural network;

segmenting the input tensor into a plurality of sub-tensors;

for each of the plurality of filters, pruning the filter, wherein the pruning comprises:

dividing a channel dimension of the filter into a plurality of channel groups, each channel group comprising a same number of two or more continuous channels of the filter, and

pruning the filter so that (1) each channel group within the filter comprises a same number of one or more non-zero weights, and (2) all non-zero weights in the filter are evenly distributed in the plurality of channel groups within the filter, wherein the pruning comprises setting one or more non-zero weights to zeros;

storing the non-zero weights in each of the plurality of channel groups of each filter as index-value pairs;

segmenting, according to the plurality of pruned channel groups, each of the plurality of filters into a plurality of sub-filters, wherein each of the plurality of sub-filters comprises the same number of non-zero weights and the corresponding index-value pair consumes a same number of bits; and

assigning the plurality of sub-tensors and the plurality of sub-filters to a plurality of processors for parallel convolution processing.

9. The system of claim 8 , wherein the operations further comprise performing the parallel convolution processing by:

iterating, by each of the plurality of processors, each non-zero weight of a sub-filter assigned to the processor; and

identifying, by the processor, a corresponding input value in a sub-tensor assigned to the processor to perform a multiply-and-accumulate (MAC) operation.

10. The system of claim 8 , wherein the operations further comprise performing the parallel convolution processing by:

for each of the non-zero weights stored as an index-value pair, identifying a corresponding input value in an assigned sub-tensor at a location identified by the channel-dimension index and the width-dimension index of the index-value pair representing each non-zero weight.

11. The system of claim 8 , wherein the segmenting each of the plurality of filters into the plurality of sub-filters according to the plurality of pruned channel groups comprises:

segmenting the filter into a plurality of chunks according to the plurality of pruned channel groups; and

segmenting each of the plurality of chunks into a plurality of horizontal planes.

12. The system of claim 11 , wherein the operations further comprise:

pruning the plurality of horizontal planes so that each of the plurality of horizontal planes comprises the same number of non-zero weights.

13. The system of claim 8 , wherein after the parallel convolution processing, the plurality of processors generate a plurality of partial sums, and

the operations further comprise:

accumulating the plurality of partial sums to obtain an output tensor; and

feeding the output tensor as an input tensor for a next layer of the neural network.

14. A non-transitory computer-readable storage medium configured with instructions executable by a plurality of processors to cause the plurality of processors to perform operations comprising:

obtaining an input tensor and a plurality of filters at a layer within a neural network;

segmenting the input tensor into a plurality of sub-tensors;

for each of the plurality of filters, pruning the filter, wherein the pruning comprises:

dividing a channel dimension of the filter into a plurality of channel groups, each channel group comprising a same number of two or more continuous channels of the filter, and

pruning the filter so that (1) each channel group within the filter comprises a same number of one or more non-zero weights, and (2) all non-zero weights in the filter are evenly distributed in the plurality of channel groups within the filter, wherein the pruning comprises setting one or more non-zero weights to zeros;

storing the non-zero weights in each of the plurality of channel groups of each filter as index-value pairs;

segmenting, according to the plurality of pruned channel groups, each of the plurality of filters into a plurality of sub-filters, wherein each of the plurality of sub-filters comprises the same number of non-zero weights and the corresponding index-value pair consumes a same number of bits; and

assigning the plurality of sub-tensors and the plurality of sub-filters to a plurality of processors for parallel convolution processing.

15. The non-transitory computer-readable storage medium of claim 14 , wherein the operations further comprise performing the parallel convolution processing by:

iterating, by each of the plurality of processors, each non-zero weight of a sub-filter assigned to the processor; and

identifying, by the processor, a corresponding input value in a sub-tensor assigned to the processor to perform a multiply-and-accumulate (MAC) operation.

16. The non-transitory computer-readable storage medium of claim 14 , wherein the segmenting each of the plurality of filters into the plurality of sub-filters according to the plurality of pruned channel groups comprises:

segmenting the filter into a plurality of chunks according to the plurality of pruned channel groups; and

segmenting each of the plurality of chunks into a plurality of horizontal planes.

17. The non-transitory computer-readable storage medium of claim 16 , wherein the operations further comprise:

pruning the plurality of horizontal planes so that each of the plurality of horizontal planes comprises the same number of non-zero weights.