时间戳

首页 > 解决方案 > numba 可以加速已经矢量化的 numpy 函数吗？

python-3.x - numba 可以加速已经矢量化的 numpy 函数吗？

问题描述

我想使用 numba 和 cuda 来并行化网格上一些厄密矩阵的对角化。让 H 是我的尺寸为 1000 x 1000 x 64 x 64 的矩阵。因此，我想获得一百万个 64 x 64 矩阵的特征向量和值。linalg 包中的 Numpy 的 eigh 函数已经做得很好了。如果我称之为linalg.eigh(H)它已经矢量化并分布在我计算机的核心上。

在 nVidia 网站上，他们展示了这样一个示例：

import numpy as np
from numba import vectorize

@vectorize(['float32(float32, float32)'], target='cuda')
def Add(a, b):
  return a + b

但在示例中，该函数尚未矢量化。基本上，我有两个问题：

  1. 是否可以在不使用 CUDA 的情况下进一步加速已经矢量化的函数 numba。
  2. 如果我想把它放在一些带有 numba 的 GPU 集群上，是否像他们在示例中建议的那样简单？

标签： python-3.xnumpynumbaeigenvector

解决方案

是否可以在不使用 CUDA 的情况下进一步加速已经矢量化的函数 numba。

不。

而且你提到的函数linalg.eigh，不是“矢量化的”（在 numpy 或 Numba 这个词的意义上），甚至不是用 Python 编写的。根据文档，它只是 Lapack 例程syevd和的包装器heevd，它将是外部库中预编译的 C 或 Fortran 代码。您看到它使用多个内核的唯一原因是因为您的 numpy 版本是针对 Lapack 的多线程实现构建的（如果您使用的是 Anaconda 版本的 numpy，则可能是Intel MKL ）。标准的 numpy 将被编译并链接到系统 Lapack，这通常是netlib 实现，并且它不是多线程的，并且不会像您描述的那样出现“矢量化”。

如果我想把它放在一些带有 numba 的 GPU 集群上，是否像他们在示例中建议的那样简单？

不，这并不容易。是不可能的。

推荐阅读