为何np.gradient远快于手写实现？解析numpy梯度函数高效原理

我测试了numpy内置np.gradient函数与采用相同中心差分法的手写梯度计算函数的耗时，通过cProfile模块发现np.gradient的耗时显著更低。这表明我的Python梯度实现方式过于基础，我希望了解np.gradient的实现原理，以掌握Python中数学运算的正确实现方法（我知晓手写函数未完全实现np.gradient的梯度计算，因为未包含边界值，但这不影响耗时对比结果）。

测试代码 #

import numpy as np
import cProfile

# 创建ndarray
N = 1000
f = np.empty((N,N))

# 测试numpy内置梯度函数的包装函数
def test_grad_np(T, n):
   for s in range(n):
       dTdx, dTdy = np.gradient(T)

# 手写梯度函数
def hc_grad(T):
   Ny, Nx = np.shape(T)
   dTdx = np.zeros((Ny, Nx))
   dTdy = np.zeros((Ny, Nx))
   for j in range(1,Nx-1):
       for i in range(1,Nx-1):
           dTdx[j,i] = (T[j,i+1] - T[j,i-1])/2.
           dTdy[j,i] = (T[j+1,i] - T[j-1,i])/2.
   return dTdx, dTdy

# 测试手写梯度函数的包装函数
def test_hc_grad(T, n):
   for s in range(n):
       dTdx, dTdy = hc_grad(T)

cProfile.run('test_hc_grad(T, 20)')
cProfile.run('test_grad_np(T, 20)')

性能测试结果 #

手写梯度函数耗时 #

144 function calls in 16.818 seconds

  Ordered by: standard name

  ncalls  tottime  percall  cumtime  percall filename:lineno(function)
      20    0.000    0.000    0.000    0.000 <__array_function__ internals>:2(shape)
       1    0.001    0.001   16.818   16.818 <string>:1(<module>)
      20    0.000    0.000    0.000    0.000 fromnumeric.py:1922(_shape_dispatcher)
      20    0.000    0.000    0.000    0.000 fromnumeric.py:1926(shape)
      20   16.803    0.840   16.804    0.840 test_speed.py:20(hc_grad)
       1    0.013    0.013   16.817   16.817 test_speed.py:30(test_hc_grad)
       1    0.000    0.000   16.818   16.818 {built-in method builtins.exec}
      20    0.000    0.000    0.000    0.000 {built-in method numpy.core._multiarray_umath.implement_array_function}
      40    0.001    0.000    0.001    0.000 {built-in method numpy.zeros}
       1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}

numpy内置梯度函数耗时 #

704 function calls (624 primitive calls) in 0.262 seconds

  Ordered by: standard name

  ncalls  tottime  percall  cumtime  percall filename:lineno(function)
      40    0.000    0.000    0.001    0.000 <__array_function__ internals>:2(empty_like)
      20    0.000    0.000    0.248    0.012 <__array_function__ internals>:2(gradient)
      40    0.000    0.000    0.001    0.000 <__array_function__ internals>:2(ndim)
       1    0.001    0.001    0.262    0.262 <string>:1(<module>)
      20    0.000    0.000    0.000    0.000 _asarray.py:110(asanyarray)
      40    0.000    0.000    0.000    0.000 _asarray.py:23(asarray)
      40    0.000    0.000    0.000    0.000 fromnumeric.py:3102(_ndim_dispatcher)
      40    0.000    0.000    0.001    0.000 fromnumeric.py:3106(ndim)
      40    0.000    0.000    0.000    0.000 function_base.py:798(_gradient_dispatcher)
      20    0.245    0.012    0.247    0.012 function_base.py:803(gradient)
      40    0.000    0.000    0.000    0.000 multiarray.py:75(empty_like)
      40    0.000    0.000    0.000    0.000 numerictypes.py:285(issubclass_)
      20    0.000    0.000    0.000    0.000 numerictypes.py:359(issubdtype)
       1    0.014    0.014    0.261    0.261 test_speed.py:15(test_grad_np)
       1    0.000    0.000    0.262    0.262 {built-in method builtins.exec}
      60    0.000    0.000    0.000    0.000 {built-in method builtins.issubclass}
      40    0.000    0.000    0.000    0.000 {built-in method builtins.len}
      60    0.000    0.000    0.000    0.000 {built-in method numpy.array}
  100/20    0.001    0.000    0.248    0.012 {built-in method numpy.core._multiarray_umath.implement_array_function}
      40    0.000    0.000    0.000    0.000 {method 'append' of 'list' objects}
       1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}

np.gradient 快的核心原因及实现原理 #

1. 避免Python级别的循环 #

你的手写函数用了两层Python原生for循环遍历数组元素，Python的循环本身开销极大——每次循环都要做类型检查、边界判断等操作，对于1000x1000的数组，要执行近百万次循环，耗时自然高。而np.gradient完全基于向量化运算，它直接对整个数组的切片进行批量计算，所有循环都在底层用C实现，避开了Python解释器的开销。

比如计算x方向梯度时，np.gradient会直接取T[:,2:] - T[:,:-2]得到所有中间点的差分结果，再除以2，整个过程是数组级别的操作，没有Python循环。

2. 内存与缓存优化 #

numpy的内置函数会利用连续内存布局和CPU缓存 locality特性。手写函数中频繁的单个元素赋值（dTdx[j,i] = ...）会导致大量零散的内存访问，而numpy的向量化操作是连续块内存读写，能充分利用CPU缓存，大幅提升效率。

3. 底层实现优化 #

np.gradient的核心逻辑在C语言实现的numpy内核中（比如_multiarray_umath模块），它还会根据数组的 dtype、维度等做针对性优化，比如使用SIMD指令集进行并行计算，进一步加速数值运算。

另外，np.gradient还处理了边界点的差分（前向/后向差分），但这部分的开销相对于向量化带来的提升可以忽略不计。

改进手写函数的思路 #

如果要提升手写代码的速度，应该抛弃Python循环，改用numpy的向量化操作：

def vectorized_grad(T):
    Ny, Nx = T.shape
    dTdx = np.zeros_like(T)
    dTdy = np.zeros_like(T)
    # 中间点用中心差分
    dTdx[:,1:-1] = (T[:,2:] - T[:,:-2])/2.
    dTdy[1:-1,:] = (T[2:,:] - T[:-2,:])/2.
    # 边界点可以补上前向/后向差分（可选）
    dTdx[:,0] = T[:,1] - T[:,0]
    dTdx[:,-1] = T[:,-1] - T[:,-2]
    dTdy[0,:] = T[1,:] - T[0,:]
    dTdy[-1,:] = T[-1,:] - T[-2,:]
    return dTdx, dTdy

这个版本的速度会和np.gradient接近，因为同样用了向量化操作。

内容的提问来源于Stack Exchange，提问作者fdv