面向规模化引力波探测的推理优化人工智能与高性能计算
摘要
我们介绍一组用于引力波探测的人工智能模型集成，这些模型在 Summit 超级计算机上使用 32 个节点（相当于 192 块 NVIDIA V100 GPU）在 2 小时内训练完成。一旦完全训练好，我们使用 NVIDIA TensorRT 对这些模型进行加速推理优化。我们将推理优化后的 AI 集成部署在 Argonne Leadership Computer Facility 的 ThetaGPU 超级计算机中以进行分布式推理。使用整个 ThetaGPU 超级计算机（由 20 个节点组成，每个节点具有 8 块 NVIDIA A100 Tensor Core GPU 和 2 块 AMD Rome CPU），我们经 NVIDIA TensorRT 优化的 AI 集成在 50 秒内处理了整月的高级 LIGO 数据（包括 Hanford 和 Livingston 数据流）。我们的推理优化 AI 集成保持了传统 AI 模型相同的灵敏度，即它识别出了该高级 LIGO 数据集中先前确认的所有已知双黑洞并合且未报告误分类，同时相比传统人工智能模型提供了 3 倍的推理加速。我们使用时间平移来量化我们的 AI 集成处理长达 5 年高级 LIGO 数据的性能。在这一合成增强的数据集中，我们的 AI 集成平均每隔一个月搜索的高级 LIGO 数据报告一次误分类。我们还给出了使用该 5 年长高级 LIGO 数据集的 AI 集成的受试者工作特征曲线。该方法提供了在规模化下进行加速的、AI 驱动的引力波探测所需的工具。
引用
@article{arxiv.2201.11133,
title = {Inference-optimized AI and high performance computing for gravitational wave detection at scale},
author = {Pranshu Chaturvedi and Asad Khan and Minyang Tian and E. A. Huerta and Huihuo Zheng},
journal= {arXiv preprint arXiv:2201.11133},
year = {2022}
}
备注
19 pages, 8 figures; v2. Accepted to Frontiers in Artificial Intelligence, Special Issue: Efficient AI in Particle Physics and Astrophysics