Mixing sparsity compression
Techniques for compressing a neural network model by mixing compression ratios (sparsity patterns) are described. The weight tensor of a neural network model is divided into weight groups. The pruning cost of compressing the weight values according to a compression ratio is determined for each weight group, and a pruning cost distribution for the compression ratio is generated from the pruning costs of the weight groups. A cost threshold can then be selected from the pruning cost distribution, and weight groups having a pruning cost below the selected cost threshold are compressed according to the compression ratio. The remaining weight groups can be compressed using one or more less aggressive compression ratios. The cost threshold can be adjusted to tune the overall sparsity and accuracy of the compressed neural network.
1 . A computer-implemented method comprising:
obtaining a weight tensor of a neural network model;
dividing the weight tensor into a plurality of weight groups;
for each of the weight groups, determining a pruning cost of compressing weight values of the weight group according to a first compression ratio that replaces a number of smallest magnitude weight values in the weight group with zeros, wherein the pruning cost is determined by calculating an average of absolute values of the smallest magnitude weight values being replaced with zeros;
generating a pruning cost distribution of the weight groups;
selecting a cost threshold from the pruning cost distribution based on a mean or median of the pruning cost distribution; and
compressing a first set of weight groups having a pruning cost below the cost threshold according to the first compression ratio to generate a sparse neural network model that occupies a smaller amount of memory than the neural network model.
2 . The computer-implemented method of claim 1 , wherein generating the sparse neural network model further includes compressing a second set of weight groups having a pruning cost at or above the cost threshold according to a second compression ratio that is less than the first compression ratio.
3 . The computer-implemented method of claim 2 , wherein the first and second compression ratios are N:4 compression ratios, wherein N being the number of original rows and 4 being the number of rows after compression, and wherein N of the first compression ratio is greater than N of the second compression ratio.
4 . The computer-implemented method of claim 1 , further comprising:
dividing remaining weight groups that are not part of the first set of weight groups into weight subgroups;
for each of the weight subgroups, determining a subgroup pruning cost of compressing the weight values of the weight subgroup according to a second compression ratio;
generating a subgroup pruning cost distribution of the weight subgroups;
selecting a subgroup cost threshold from the subgroup pruning cost distribution; and
compressing a set of weight subgroups having a subgroup pruning cost below the subgroup cost threshold according to the second compression ratio.
5 . The computer-implemented method of claim 1 , wherein the sparse neural network model includes one or more uncompressed weight groups.
6 . The computer-implemented method of claim 1 , wherein generating the sparse neural network model further includes generating metadata tags for each non-zero weight value remaining in the compressed first set of weight groups indicating a position of the non-zero weight value.
7 . The computer-implemented method of claim 1 , further comprising:
retraining the neural network using the sparse neural network model; and
determining whether the sparse neural network model satisfies a network accuracy target.
8 . The computer-implemented method of claim 7 , further comprising:
determining that the sparse neural network model fails to satisfy the network accuracy target; and
selecting a lower cost threshold from the cost threshold distribution to generate a second sparse neural network model.
9 . The computer-implemented method of claim 1 , wherein each weight group is a set of rows of weight values of a two-dimensional weight tensor.
10 . The computer-implemented method of claim 9 , wherein the two-dimensional weight tensor is obtained by flattening weight matrices of a multi-layer neural network model.
11 . The computer-implemented method of claim 1 , further comprising constraining weight groups belonging to a same layer of the neural network model to use a same compression ratio.
12 . The computer-implemented method of claim 1 , wherein the first compression ratio is a L:R compression ratio, wherein L represents the number of original weight values, and R represents the number of weight values after pruning, and wherein R is less than or equal to 4.
13 . The computer-implemented method of claim 12 , wherein L is less than or equal to 16.
14 . The computer-implemented method of claim 1 , wherein the first compression ratio is selected based on a native compression support of a systolic array.
15 . The computer-implemented method of claim 14 , wherein the native compression support of the systolic array is determined by the number of feature map channels provided to a row group of the systolic array.
16 . A non-transitory computer readable medium having stored therein instructions that, when executed by one or more processors, cause the one or more processors to performing operations including:
obtaining a weight tensor of a neural network model;
selecting a first cost threshold associated with a first compression ratio;
compressing a first set of weight groups having a pruning cost below the first cost threshold according to the first compression ratio, wherein the pruning cost is determined by calculating an average of absolute values of the smallest magnitude weight values being replaced with zeros, and the first cost threshold is selected based on a mean or median of a pruning cost distribution of the first set of weight groups; and
compressing a second set of weight groups according to a second compression ratio that is less than the first compression ratio to generate a sparse neural network model that occupies a smaller amount of memory than the neural network model.
17 . The non-transitory computer readable medium of claim 16 , wherein the second set of weight groups has a pruning cost below a second cost threshold associated with the second compression ratio.
18 . The non-transitory computer readable medium of claim 16 , wherein the operations further include:
generating metadata tags for each non-zero weight value remaining in the compressed first set of weight groups and the compressed second set of weight groups indicating a position of the non-zero weight value.
19 . The non-transitory computer readable medium of claim 16 , wherein the instructions are compiler instructions for a compiler executed by the one or more processors.
20 . The non-transitory computer readable medium of claim 16 , wherein each weight group is a set of rows of weight values of a two-dimensional weight tensor.