机器学习驱动的混合分段多方与表格状态方程中的守恒量到原变量转换

我们提出了一种新颖的机器学习（ML）方法，以加速守恒量到原变量的反演，重点关注混合分段多方和表格状态方程。传统的求根技术计算成本高昂，尤其对于大规模相对论流体动力学模拟而言。为解决此问题，我们采用前馈神经网络（NNC2PS 和 NNC2PL），在 PyTorch 中训练，并使用 NVIDIA TensorRT 针对 GPU 推理进行优化，在实现显著加速的同时将精度损失降至最低。NNC2PS 模型的 $L_1$ 和 $L_\infty$ 误差分别为 $4.54 \times 10^{-7}$ 和 $3.44 \times 10^{-6}$，而 NNC2PL 模型则表现出更低的误差值。与传统的根查找方法相比，采用混合精度部署的 TensorRT 优化显著提升了性能。具体而言，对于包含 1,000,000 个点的数据集，NNC2PS 的混合精度 TensorRT 引擎的推理速度比传统的单线程 CPU 实现快约 400 倍。在 Delta 超级计算机（双 AMD 64 核 2.45 GHz Milan 处理器；8 个配备 40 GB HBM2 RAM 和 NVLink 的 NVIDIA A100 GPU）的整个计算节点上进行理想并行化处理，预测在处理 800 万个数据点时，TensorRT 相对于最优并行化的数值方法可实现 25 倍的加速。此外，该机器学习方法随着数据集规模的增加表现出次线性扩展。我们发布了所开发的科学软件，以便进一步验证和扩展我们的发现。这项工作强调了机器学习与 GPU 优化和模型量化相结合，在加速相对论流体动力学模拟中的守恒量到原变量反演方面的潜力。