如何在 Python 中使用 M2 max GPU 进行多处理？

Question

我正在编写物理模拟代码 - 我将省略技术细节，但本质上我正在多次运行我的模拟（由许多大型数学运算组成）以获得不同输入参数的列表：即

runSimulation(a,b,c,...)

。在伪代码中，这是一个概述：

paramArray = [[a1,b1,c1,...];[a2,b2,c2,...];...;[a_n,b_n,c_n]]
for param in paramArray:
    runSimulation(param)

runSimulation()

代码本身大约有 800 行代码并且非常复杂，因此除非有要求，否则我将省略任何额外的细节。

我需要运行很多很多参数组合（>250），所以我正在寻找最优化（最快）的方法来运行这些计算。每个调用都是完全独立的，因此可以并行化。我需要一些有关单个 GPU 上的多处理的帮助。 不涉及机器学习，只是非常大的物理模拟。

我最初的代码使用 CPU

multiprocessing

 模块（池为 12 个），并设法在 

150 秒左右为每个输出生成一个文件。然而，我有一个 m2 max 芯片，并认为我可以利用 GPU 来进行这些大型数学运算。

我将整个代码从

numpy

转换为

torch

，并将

runSimulation()

内的所有内容发送到涉及大矩阵乘法的GPU。现在我的代码需要 

33 小时。我知道一切都是按顺序运行的，不再是并行的。如何解决这个问题并使用单个 GPU 和 12 个核心来加速我的代码？

非常感谢！值得注意的是，我对 GPU 计算非常陌生——因此我相信我的效率很低。

Answer 1

我受益匪浅

在此输入链接描述