[论文解读] FastONN -- Python based open-source GPU implementation for Operational Neural Networks

Junaid Malik, Serkan Kıranyaz+1 more|arXiv (Cornell University)|Jun 3, 2020

Model Reduction and Neural Networks参考文献 7被引用 24

一句话总结

FastONN 是一个基于 Python 的开源 GPU 库，用于训练操作神经网络（ONNs），通过引入操作神经元的向量化公式，实现高效、灵活且可微分的计算。它利用自动反向模式微分进行反向传播，支持可自定义的算子集合与梯度流，并包含用于性能跟踪、检查点保存和指标监控的辅助模块，显著加速了 GPU 硬件上的 ONN 训练。

ABSTRACT

Operational Neural Networks (ONNs) have recently been proposed as a special class of artificial neural networks for grid structured data. They enable heterogenous non-linear operations to generalize the widely adopted convolution-based neuron model. This work introduces a fast GPU-enabled library for training operational neural networks, FastONN, which is based on a novel vectorized formulation of the operational neurons. Leveraging on automatic reverse-mode differentiation for backpropagation, FastONN enables increased flexibility with the incorporation of new operator sets and customized gradient flows. Additionally, bundled auxiliary modules offer interfaces for performance tracking and checkpointing across different data partitions and customized metrics.

研究动机与目标

• 为操作神经网络（ONNs）缺乏高效、可扩展且 GPU 加速的训练框架提供解决方案。
• 在可微分神经网络框架中实现异构非线性算子的灵活集成。
• 为多样化 ONN 应用支持可自定义的梯度流和算子集合。
• 提供模块化工具用于训练过程中的性能跟踪、检查点保存和指标监控。
• 通过 GPU 优化的向量化计算和自动微分，加速 ONN 训练。

提出的方法

• 提出一种新颖的操作神经元向量化公式，以实现高效的 GPU 计算。
• 采用自动反向模式微分，实现 ONNs 中端到端的反向传播。
• 设计模块化库架构，以支持插件式算子集合与自定义梯度流。
• 使用类似 PyTorch 的计算图构建方式，实现与现有深度学习工作流的无缝集成。
• 集成辅助模块，用于日志记录、跨数据分区的检查点保存以及自定义指标计算。
• 通过使用 CUDA 加速的张量运算，实现性能优化。

实验结果

研究问题

• RQ1 如何通过 GPU 加速高效训练操作神经网络？
• RQ2 向量化且可微分的公式对 ONN 训练速度和灵活性有何影响？
• RQ3 模块化开源库能否在 ONNs 中支持多样化的算子集合与自定义梯度流？
• RQ4 性能跟踪与检查点集成如何提升训练的稳定性和可扩展性？
• RQ5 自动微分在多大程度上可简化复杂 ONN 层的实现？

主要发现

• FastONN 通过 GPU 加速的向量化计算实现了操作神经元的显著训练加速。
• 该库通过自动反向模式微分支持新算子集合与自定义梯度流的灵活集成。
• 用于性能跟踪、检查点保存和自定义指标的辅助模块提升了训练的可重现性与可扩展性。
• 向量化公式降低了计算开销，并提高了 GPU 硬件上的内存效率。
• 开源特性与基于 Python 的设计促进了可扩展性与社区贡献。
• 该框架与现有深度学习工作流表现出良好的兼容性，尤其适用于依赖自动微分的场景。