梯度函数求值程序性能优化求助：SymPy替代方案及提速方法

问题描述 #

我编写了一段代码，用于让点沿梯度方向移动以寻找3D函数的最小值，但运行速度极慢。核心逻辑在roll.roll()中，通过循环重复计算函数值、沿梯度移动点实现。我怀疑慢的原因是calculate()函数效率低下，或是SymPy的符号操作拖慢了速度。

另外我需要点在函数值较高的位置移动更快（不只是斜率大的区域），这也增加了实现难度。

原代码如下：

import sympy as smp

x, y = smp.symbols('x y')
eq = 1*smp.exp(-((x-5)/5)**2 - ((y-1)/2)**2) + \
    2*smp.exp(-((x+3)/2)**2 - ((y-3)/2)**2) + \
    3*smp.exp(-((x-4)/2)**2 - ((y-7)/2)**2)

# Evaluates the 2 input sympy symbolic function "expression" at points (x1,y1)
def calculate(expression,x1,y1):
    EQ = smp.lambdify((x,y), expression, 'numpy')
    return EQ(x1,y1)

class roll:
    xDiff = smp.diff(eq,x)
    yDiff = smp.diff(eq,y)
    normalize = eq/smp.sqrt(xDiff**2 + yDiff**2)

    def roll(x,y,duration):
        (x,y) = (x,y)
        for i in range(0,duration):
            (x,y) = (
                x-calculate((roll.normalize*roll.xDiff),x,y),
                y-calculate((roll.normalize*roll.yDiff),x,y)
                )
        return (x,y)

print(roll.roll(1,2,10))

核心优化方案 #

1. 避免重复调用lambdify #

你的calculate()函数每次循环都会重新将SymPy表达式编译为NumPy可调用函数，这是最大的性能瓶颈。应提前预编译所有需要用到的表达式，而非每次循环重复编译：

import sympy as smp

x, y = smp.symbols('x y')
eq = 1*smp.exp(-((x-5)/5)**2 - ((y-1)/2)**2) + \
    2*smp.exp(-((x+3)/2)**2 - ((y-3)/2)**2) + \
    3*smp.exp(-((x-4)/2)**2 - ((y-7)/2)**2)

# 提前计算所有需要的导数和步长表达式
xDiff = smp.diff(eq, x)
yDiff = smp.diff(eq, y)
normalize = eq / smp.sqrt(xDiff**2 + yDiff**2)
step_x_expr = normalize * xDiff
step_y_expr = normalize * yDiff

# 预编译为NumPy函数，仅执行一次
step_x = smp.lambdify((x, y), step_x_expr, 'numpy')
step_y = smp.lambdify((x, y), step_y_expr, 'numpy')

class roll:
    @staticmethod
    def roll(x, y, duration):
        for _ in range(duration):
            x -= step_x(x, y)
            y -= step_y(x, y)
        return (x, y)

print(roll.roll(1, 2, 10))

2. 用Autograd替代SymPy做梯度计算 #

Autograd支持自动计算NumPy函数的梯度，采用数值自动微分，速度远快于SymPy的符号微分，更适合迭代场景：

import autograd.numpy as np
from autograd import grad

def eq_np(x):
    return (1*np.exp(-((x[0]-5)/5)**2 - ((x[1]-1)/2)**2) +
            2*np.exp(-((x[0]+3)/2)**2 - ((x[1]-3)/2)**2) +
            3*np.exp(-((x[0]-4)/2)**2 - ((x[1]-7)/2)**2))

# 自动生成梯度函数
grad_eq = grad(eq_np)

def roll(x, y, duration):
    pos = np.array([x, y])
    for _ in range(duration):
        g = grad_eq(pos)
        # 实现高函数值位置移动更快的逻辑：函数值 * 归一化梯度
        f_val = eq_np(pos)
        norm_g = np.linalg.norm(g)
        step = f_val * g / norm_g
        pos -= step
    return pos

print(roll(1, 2, 10))

3. 用SciPy优化器直接实现 #

如果不需要手动实现梯度下降逻辑，SciPy的scipy.optimize.minimize底层采用C实现，效率远高于自定义代码：

import numpy as np
from scipy.optimize import minimize

def eq_np(x):
    return (1*np.exp(-((x[0]-5)/5)**2 - ((x[1]-1)/2)**2) +
            2*np.exp(-((x[0]+3)/2)**2 - ((x[1]-3)/2)**2) +
            3*np.exp(-((x[0]-4)/2)**2 - ((x[1]-7)/2)**2))

# 初始点(1,2)，使用L-BFGS-B高效寻优
result = minimize(eq_np, [1, 2], method='L-BFGS-B')
print(result.x)

若需要自定义步长逻辑，可在梯度函数中嵌入函数值权重的计算规则。

总结 #

• 优先级最高的优化是提前预编译SymPy表达式，消除重复编译的开销；
• 追求极致速度时，优先选择Autograd（自定义梯度下降）或SciPy优化器（无需手动实现）替代SymPy的符号计算；
• 高函数值位置移动更快的需求，可通过函数值 × 归一化梯度的步长公式实现，上述方案均支持该逻辑。

内容的提问来源于stack exchange，提问作者Ben