Array Languages Make Neural Networks Fast
Artjoms Šinkarovs, H.-N. Vießmann, and S.-B. Scholz, “Array languages make neural networks fast,” in Proceedings of the 7th ACM SIGPLAN international workshop on libraries, languages and compilers for array programming, Jun. 2021, pp. 39–50. doi: 10.1145/3460944.3464312.
Abstract
Most implementations of machine learning algorithms are based on special-purpose frameworks such as TensorFlow or PyTorch. While these frameworks are convenient to use, they introduce multi-million lines of code dependency that one has to trust, understand and potentially modify. As an alternative, this paper investigates a direct implementation of a state of the art Convolutional Neural Network (CNN) in an array language. While our implementation requires 150 lines of code to define the special-purpose operators needed for CNNs, which are readily provided through frameworks such as TensorFlow and PyTorch, our implementation outperforms these frameworks by factors 2 and 3 on a fixed set of hardware — a 64-core GPU-accelerated machine; for a simple example network. The resulting specification is written in a rank-polymorphic data-parallel style, and it can be immediately leveraged by optimising compilers. Indeed, array languages make neural networks fast.
Additional Information
This talk is recorded and is available here. The code is available here.