PCTC: Hardware and Software Co-design for Pruned Capsule Networks on Tensor Cores

Capsule Networks (CapsNets) are a generation of image classifiers that have taken the spotlight compared to convolutional neural networks (CNNs). Unlike CNNs, CapsNets are robust to affine transformation and are able to learn spatial relationships between features of an image. Since CapsNets require significant computing horsepower due to intensive matrix operations, GPUs have become the primary hardware platforms for the execution of CapsNets. In particular, GPUs equipped with tensor cores (TCs) are an attractive solution to address the computational requirements of CapsNets as TCs are designed to accelerate matrix operations. However, CapsNets deployed into TCs underutilize computational units as TCs are designed to run dense matrix operations. We propose pruned capsule TC (PCTC) which is a software/hardware co-design approach and avoids underutilization of TC resources. In particular, PCTC changes the sequence of matrix operations for capsule layers so that sparse operations are eliminated. PCTC further enhances the execution of CapsNets on TCs by eliminating those matrix operations that are not necessary to maintain the accuracy of the network. Quite often, CapsNets are designed with large capsules to increase accuracy. By pruning individual capsules, it is feasible to reduce the over-provisioned parameter space and reduce energy consumption in CapsNets. Evaluation results reveal that PCTC can achieve 31% energy saving for CapsNet inference, with negligible accuracy loss.