Deep Quantization of Graph Neural Networks with Run-Time Hardware-Aware Training

Abstract

In this paper, we investigate the benefits of hardware-aware quantization in the gFADES hardware accelerator targeting Graph Convolutional Networks (GCNs). GCNs are a type of Graph Neural Networks (GNNs) that combine sparse and dense data compute requirements that are challenging to meet in resource-constrained embedded hardware. The gFADES architecture is optimized to work with the pruned data representations typically present in graph neural networks for the graph structure and features. It is described in High-Level Synthesis (HLS) which enables efficient design-space exploration of mixed precision hardware configurations. In this work, the mixed-precision design is embedded in the forward pass of the PyTorch back-propagation training loop to enable run-time hardware-aware training. It uses different data types to represent adjacency, feature, weight, internal, and output values which allows for a fine-grained optimization at the tensor level. The resulting hardware configuration after training reduces precision to a 4-bit data type for all inputs. It achieves little to no degradation in the classification accuracy, when training on the Planetoid database dataset, compared to the original 32-bit floating-point. The optimized hardware design running on an AMD/Xilinx Zynq Ultrascale+ FPGA device achieves over $600\times$ speedup compared to the optimized PyTorch software implementation running on the multi-core ARM CPU in the processing system.