如何让Python代码在GPU上运行 – 搞英语 → 看世界

作为一名软件开发人员，我希望能够指定某些代码在 GPU 内运行，以便它可以并行执行。具体而言，本文演示了如何使用 Python 3.9 在配备 Apple M1 Pro 芯片的 MacBook Pro 的 GPU 上运行代码。

适合 GPU 的任务如下：

汇总数组中的值（map/reduce）
矩阵乘法、数组运算
图像处理（图像是像素阵列）
使用上述组合的机器学习

为了使用 GPU，我选择渲染Mandelbrot 集。这篇文章还将比较 MacBook Pro 的 CPU 与 GPU 的性能。 github 上提供了该项目的完整代码，因此您可以自己尝试一下。

编写在 GPU 上运行的代码：

在 Python 中，通过 GPU 运行代码并不是本机功能。一个流行的库是TensorFlow 2.14 ，截至 2023 年 10 月，它可与 MacBook Pro M1 GPU 硬件配合使用。尽管 TensorFlow 是为机器学习而设计的，但它提供了一些利用 GPU 并行化的基本数组操作函数。要使 GPU 工作，您需要安装 Apple 提供的TensorFlow-Metal 包。如果没有它，即使使用 TensorFlow 包，您也只能停留在 CPU 领域。

与传统的“程序逻辑”相比，在 TensorFlow（以及一般的 GPU 库）中进行编程需要进行一些不同的思考。 TensorFlow 不是一次处理一个单元，而是同时处理所有元素。数据列表需要保存在特殊的 Tensor 对象中（接受 numpy 数组作为输入）。加法、减法和乘法等运算在张量上超载。在幕后，当您添加/减去/相乘张量时，它会将数据分解为更小的块，并将工作并行地分配给 GPU。不过，执行此操作会产生开销，而 CPU 首当其冲。如果您的数据集很小，GPU 方法实际上运行速度会更慢。随着数据集的增长，GPU 最终将被证明更加高效，并使以前仅使用 CPU 无法完成的任务成为可能。

您如何知道您的 GPU 正在被使用？

要查看 CPU 和 GPU 使用情况，请打开“活动监视器”，然后打开“窗口”->“GPU 历史记录”（命令 4），然后打开“窗口”->“CPU 历史记录”（命令 3）。

运行TensorFlow-Metal 指令第 4 步中的脚本，该脚本会启动一堆张量并使用测试数据构建基本的机器学习模型。

在你的 GPU 历史记录窗口中，你应该会看到它已经达到最大值，如下所示：

M1 Pro GPU 活动

曼德尔布罗特代码：

Mandelbrot 集是 1978 年的一个奇怪的数学发现。维基文章对它的工作原理有很好的描述。基本上，它涉及检查笛卡尔坐标系中的每个点，以查看该点的值在输入“简单”方程时是否稳定或发散到无穷大。它恰好涉及复数（具有虚数部分，Y 值提供该部分），但 Python 代码可以很好地处理它。当你绘制它时，你得到的是一个美丽/怪异的图像，本质上是分形的。你可以继续放大它的某些部分，它会显示隐藏在较小视图中的较大视图的分形表示，一直到计算机可以接受的程度。

Mandelbrot 集的完整视图，由本项目中的代码生成：

曼德尔布罗特大视野

这是构建曼德尔布罗特集的简单“程序”方法。请注意，它会逐一计算每个像素。

 def mandelbrot_score(self, c: 复杂, max_iterations: int) -> float: 
  
    ”“” 
  
    计算所提供的给定复数的曼德尔布罗分数。 
  
    曼德尔布罗网格中的每个像素都有由 x + 1j*y 确定的 ac 值（1j 是 sqrt(-1) 的表示法） 
  
 
  
    :param c: 要测试的复数 
  
    :param max_iterations: 压缩 z 值的次数 (z ** 2 + c) 
  
    :return: 1 如果 c 值稳定，或者值 0 >= x > 1 表示它的分歧速度 
  
            （较低意味着它背离得更快）。 
  
    ”“” 
  
    z = 0 
  
    对于范围内的 i（最大迭代次数）： 
  
        z = z ** 2 + c 
  
        如果abs(z) > 4： 
  
            # 当它超过abs > 4后，假设它会趋于无穷大 
  
            # 返回相对于 max_iterations 多久开始出现峰值 
  
            返回 i / 最大迭代次数 
  
 
  
    # c值稳定 
  
    返回1 
  
 
  
# 下面是存储库的 MandelbrotCPUBasic 类中使用的逻辑的简化版本： 
  
 
  
# 设置一个 numpy 像素数组网格 
  
像素 = np.zeros((500, 500)) 
  
 
  
# 计算每个像素的散度值 
  
对于范围 (500) 内的 y： 
  
    对于范围 (500) 内的 x： 
  
        # 计算该像素的“常数” 
  
        c = x + 1j*y 
  
 
  
        # 获取该像素的散度分数 
  
        分数 = mandelbrot_score(c, 50) 
  
 
  
        # 将分数保存在像素网格中 
  
        像素[y][x] = 分数

这是 TensorFlow 2.x 的实现方法。请注意，它在tensor_flow_step函数的第一行中同时对所有值进行操作，并将输入值返回到调用循环。

 def tensor_flow_step(self, c_vals_, z_vals_, divergence_scores_): 
  
    ”“” 
  
    compute_mandelbrot_tensor_flow()的处理步骤， 
  
    一次计算所有像素。 
  
 
  
    ：param c_vals_：每个坐标的复数值数组 
  
    :param z_vals_: 每个坐标的 z 值，从 0 开始，每一步重新计算 
  
    :param divergence_scores_: 每个像素发散之前所进行的迭代次数 
  
    :return: 更新后的输入 
  
    ”“” 
  
 
  
    z_vals_ = z_vals_*z_vals_ + c_vals_ 
  
 
  
    # 查找未发散的 z 值，并仅增加这些元素 
  
    not_diverged = tf.abs(z_vals_)

结果：

以下是使用 CPU 与 GPU 通过 TensorFlow 生成不同大小的 Mandelbrot 图像的结果。请注意，TensorFlow 代码完全相同，我只是使用with tf.device()方法强制它使用 CPU/GPU。

在不同分辨率下生成 Mandelbrot 的时间 CPU 与 GPU

在 TensorFlow GPU 和 CPU 之间，我们可以看到它们在 5000 x 5000 之前大致相同。然后在 10000 x 10000 时 GPU 稍微领先。在 15000 x 15000 下，GPU 的速度几乎是原来的两倍！这显示了从 CPU 到 GPU 的资源编组如何增加开销，但一旦数据集的大小足够大，任务的数据处理方面就会超过使用 GPU 的额外成本。

有关这些结果的详细信息：

日期：2023年10月29日
MacBook Pro（16 英寸，2021 年）
芯片：苹果M1 Pro
内存：16GB
macOS 12.7
Python 3.9.9
numpy 1.24.3
张量流2.14.0
张量流金属 1.1.0

算法/设备类型	图片大小	时间（秒）
中央处理器基础	500×500	0.484236
中央处理器基础	2500×2500	12.377721
中央处理器基础	5000×5000	47.234169
TensorFlow GPU	500×500	0.372497
TensorFlow GPU	2500×2500	2.682249
TensorFlow GPU	5000×5000	13.176994
TensorFlow GPU	10000×10000	42.316472
TensorFlow GPU	15000×15000	170.987643
TensorFlow CPU	500×500	0.265922
TensorFlow CPU	2500×2500	2.552139
TensorFlow CPU	5000×5000	12.820812
TensorFlow CPU	10000×10000	46.460504
TensorFlow CPU	15000×15000	328.967006

注意：对于 CPU Basic 算法，我在 5000 x 5000 之后放弃了，因为 10000 x 10000 图像变得非常低，并且这一点得到了很好的证明，TensorFlow 的实现要快得多。

好奇它如何在您的硬件上工作？为什么不试试呢？该项目的代码可在 github 上找到。

关于在 GPU 上运行 Python 代码的其他想法：

另一个值得一提的项目是PyOpenCL 。它封装了OpenCL ，这是一个用于编写针对不同设备（包括 GPU）执行的函数的框架。 OpenCL 需要 GPU 制造商提供的兼容驱动程序才能工作（例如 AMD、Nvidia、Intel）。

我实际上尝试让 PyOpenCL 在我的 Mac 上运行，但结果Apple 不再支持 OpenCL 。我还遇到过对CUDA 的引用，它类似于 OpenCL，更成熟一点，但它仅适用于 Nvidia GPU。如果您碰巧有 Nvidia 显卡，您可以尝试使用PyCUDA 。

CUDA 和 OpenCL 之于 GPU 并行处理就像 DirectX 和 OpenGL 之于图形处理一样。像 DirectX 这样的 CUDA 是专有的，但非常强大，而 OpenCL 和 OpenGL 本质上是“开放”的，但缺乏某些内置功能。不幸的是，在配备 M1 芯片的 MacBook Pro 上，这两个选项都不是。截至 2023 年 10 月，TensorFlow 是我能看到的唯一选择。网上有很多关于在 Mac 上使用 PyOpenCL 的过时信息，但当我尝试让它运行时，一切都陷入了死胡同。

这篇文章的灵感/来源：

https://people.duke.edu/~ccc14/sta-663/CUDAPython.html – 使用 CUDA 绘制 Mandelbrot 集的示例。
https://realpython.com/mandelbrot-set-python/ – 在 Python 中绘制 Mandelbrot 集的完整教程，包括着色和放大旋臂。
https://dzone.com/articles/mandelbrot-set-in-tensorflow – 使用 TensorFlow 的示例。

如何让 Python 代码在 GPU 上运行一文首先出现在Laurence Gellert 的博客上。

原文： https://www.laurencegellert.com/2023/10/how-make-python-code-run-on-the-gpu/