如果数组是在 GPU 中定义的，为什么 np.einsum() 可以在 GPU 上运行？

Question

我最近意识到，只要在 GPU 中定义变量（或数组），numpy 函数运行就和 cupys 一样快。我找不到一种方法来监视 numpy 函数是否实际在 GPU 上执行。如果有人有类似的经历，请分享您的答案。谢谢。

我正在使用 google colabpro GPU 运行time.

下面是两组代码：

import numpy as np
import cupy as cp

第 1 组：

a = np.random.randn(500, 500, 500)
b = np.random.randn(500, 500, 500)   
start_time = time.time()
for i in range(1): 
  c = np.einsum('ijk,ikm->ijm', a, b)        
end_time = time.time()
print('forwar gpu time')
print(end_time - start_time)

forwar gpu time
55.88586902618408

第 2 组：

a = cp.random.randn(500, 500, 500)   # change to cupy
b = cp.random.randn(500, 500, 500)   # change to cupy
start_time = time.time()
for i in range(1): 
  c = np.einsum('ijk,ikm->ijm', a, b)    # remain numpy    
end_time = time.time()
print('forwar gpu time')
print(end_time - start_time)

forwar gpu time
0.0009937286376953125

Answer 1

让我们看一下 einsumfunc.py 中的 numpy.einsum:

@array_function_dispatch(_einsum_dispatcher, module='numpy')
def einsum(*operands, out=None, optimize=False, **kwargs):

嗯，那个装饰师看起来很有前途；您确实在描述参数类型的函数 dispatching。让我们看一下 overrides.py

中的 array_function_dispatch

def array_function_dispatch(dispatcher, module=None, verify=True,
                            docs_from_dispatcher=False):
    """Decorator for adding dispatch with the __array_function__ protocol.
    See NEP-18 for example usage.

这将我们带到更具可读性的 NEP-18：“我们提出 __array_function__ 协议，以允许 NumPy 函数的参数定义该函数如何对它们进行操作...”所以NumPy 函数检查其参数的 __array_function__ 或 __array_ufunc__ (NEP-13)。

CuPy 在 core.pyx（一个 Cython 文件）中定义了 cupy.ndarray.__array_function__。它在类似组织的 CuPy 模块中查找 NumPy 函数的名称，并根据参数调用它。因此，当您调用 np.einsum(...) 时，它最终还是找到了 cp.einsum(...)：

cdef class ndarray:
...
...
...
    def __array_function__(self, func, types, args, kwargs):
        try:
            module = functools.reduce(
                getattr, func.__module__.split('.')[1:], cupy)
            cupy_func = getattr(module, func.__name__)
        except AttributeError:
            return NotImplemented
        if cupy_func is func:
            # avoid NumPy func
            return NotImplemented
        for t in types:
            if t not in _HANDLED_TYPES:
                return NotImplemented
        return cupy_func(*args, **kwargs)

如果数组是在 GPU 中定义的，为什么 np.einsum() 可以在 GPU 上运行？

why np.einsum() can run on GPU if the arrays are defined in GPU?

python

gpu

google-colaboratory

如果数组是在 GPU 中定义的，为什么 np.einsum() 可以在 GPU 上 运行？

why np.einsum() can run on GPU if the arrays are defined in GPU?

python

gpu

google-colaboratory

如果数组是在 GPU 中定义的，为什么 np.einsum() 可以在 GPU 上运行？