We introduce a library called Push that takes a probabilistic programming approach to Bayesian deep learning (BDL). This library enables concurrent execution of BDL inference algorithms on multi-GPU hardware for neural network (NN) models. To accomplish this, Push introduces an abstraction that represents an input NN as a particle. Push enables easy creation of particles so that an input NN can be replicated and particles can communicate asynchronously so that a variety of parameter updates can be expressed, including common BDL algorithms. Our hope is that Push lowers the barrier to experimenting with BDL by streamlining the scaling of particles across GPUs. We evaluate the scaling behavior of particles on single-node multi-GPU devices on vision and scientific machine learning (SciML) tasks.
翻译:我们介绍一个名为Push的库,它采用概率编程方法实现贝叶斯深度学习。该库支持在面向神经网络模型的多GPU硬件上并发执行BDL推理算法。为此,Push引入了一种将输入神经网络抽象为粒子的表示方法。通过简化粒子创建以实现输入NN的复制,并支持粒子间的异步通信,该库可表达多种参数更新方式(包括常见BDL算法)。我们期望Push能通过简化跨GPU粒子的扩展过程,降低BDL实验的门槛。我们在单节点多GPU设备上,针对视觉任务与科学机器学习任务评估了粒子的扩展性能。