通过混合精度加速科学机器学习并降低内存使用
机器学习
2024-01-31 v1
摘要
科学机器学习(SciML)已成为一种灵活的方法，用于解决复杂的计算科学和工程问题。在该领域，物理信息神经网络(PINNs)和深操作网络(DeepONets)是解决部分微分方程的领先技术，通过融合物理方程和实验数据。然而，训练PINNs和DeepONets需要大量计算资源，包括长时间的计算和大量内存。在寻求计算效率的过程中，使用半精度(float16)而非常规的单精度(float32)或双精度(float64)进行神经网络训练已获得广泛关注，由于其减少计算时间和内存消耗的内在优势。然而，我们发现float16不能应用于SciML方法，因为在训练初期会出现梯度发散、权重更新为零以及无法收敛到局部最小值。为克服这些限制，我们探索了混合精度，这是一种结合float16和float32数值格式的技术，用于降低内存使用和提高计算速度。我们的实验表明，混合精度训练不仅显著降低了训练时间和内存需求，还保持了模型准确性。我们还通过理论分析强化了我们的经验观察。该研究对SciML在各种计算应用中具有广泛意义。
引用
@article{arxiv.2401.16645,
title = {Speeding up and reducing memory usage for scientific machine learning via mixed precision},
author = {Joel Hayford and Jacob Goldman-Wetzler and Eric Wang and Lu Lu},
journal= {arXiv preprint arXiv:2401.16645},
year = {2024}
}
备注
25 pages, 7 figures