如何加速Python中均方误差(MSE)优化的计算速度？

问题背景 #

论文研究中需快速完成带L1正则的MSE优化任务，当前用Python 3.9（Mac环境）scipy.optimize.minimize（带Jacobian）处理小规模数据耗时约40ms，但实际数据规模更大，且需重复执行最多50次。Mac下Python多进程启动慢、fork不稳定，希望优化单任务速度替代多进程。曾尝试CVXPY实现，结果与scipy不符且速度更慢，寻求免费开源求解器的加速方案。

当前实现代码 #

目标函数 #

def _residuals_mse(coef, unshimmed_vec, coil_mat, factor, reg_factor, reg_factor_channel):
    """ Objective function to minimize the mean squared error (MSE)

    Args:
        coef (numpy.ndarray): 1D array of channel coefficients
        unshimmed_vec (numpy.ndarray): 1D flattened array (point) 
        coil_mat (numpy.ndarray): 2D flattened array (point, channel) of masked coils
                                  (axis 0 must align with unshimmed_vec)
        factor (float): Scale the output to avoid positive directional linesearch
        reg_factor (float): Regularization strength
        reg_factor_channel (numpy.ndarray): Channel-wise regularization scaling

    Returns:
        scalar: Residual for least squares optimization 
    """
    residual = unshimmed_vec + coil_mat @ coef
    return np.mean(residual ** 2) / factor + reg_factor * np.mean(np.abs(coef) / reg_factor_channel)

注：原代码中self相关参数改为显式传入，避免闭包开销；用矩阵乘法替代np.sum(coil_mat * coef, axis=1)，提升运算效率

Jacobian函数 #

def _residuals_mse_jacobian(coef, unshimmed_vec, coil_mat, factor, reg_factor, reg_factor_channel):
    """ Jacobian of the MSE regularization objective

    Args:
        coef (numpy.ndarray): 1D array of channel coefficients
        unshimmed_vec (numpy.ndarray): 1D flattened array (point) of the masked unshimmed map
        coil_mat (numpy.ndarray): 2D flattened array (point, channel) of masked coils
                                  (axis 0 must align with unshimmed_vec)
        factor (float): Scaling factor
        reg_factor (float): Regularization strength
        reg_factor_channel (numpy.ndarray): Channel-wise regularization scaling

    Returns:
        jacobian (numpy.ndarray) : 1D array of the gradient
    """
    n_points = unshimmed_vec.size
    b = 2 / (n_points * factor)
    residual = unshimmed_vec + coil_mat @ coef
    # 向量化计算MSE项梯度
    mse_grad = b * coil_mat.T @ residual
    # L1正则项梯度
    l1_grad = reg_factor * np.sign(coef) / reg_factor_channel / coef.size
    return mse_grad + l1_grad

注：消除原代码中的Python循环，替换为Numpy向量化运算；修正原代码中self.b的变量错误

主程序 #

import numpy as np 
import scipy.optimize as opt
from numpy.random import default_rng

rand = default_rng(seed=0)
reg_factor_channel = rand.integers(1, 10, size=9) 
coef_init = np.zeros(9)
unshimmed_vec = np.random.randint(100, size=(150))
coil_mat = np.random.randint(100, size=(150,9))
factor = 2 
reg_factor = 5 

currents_sp = opt.minimize(
    _residuals_mse,
    coef_init,
    args=(unshimmed_vec, coil_mat, factor, reg_factor, reg_factor_channel),
    method='L-BFGS-B',
    jac=_residuals_mse_jacobian,
    options={'maxiter': 1000, 'disp': False}
)

注：将优化器更换为L-BFGS-B，更适合带正则项的光滑目标优化

CVXPY尝试问题说明 #

曾用CVXPY实现但结果不符、速度更慢，原因主要是：

• 正则项计算逻辑与scipy实现不一致：原代码中x3 = self.reg_factor / self.reg_factor_channel@ cp.abs(coef) / m的矩阵乘法顺序、缩放比例存在偏差
• SCS求解器在小规模问题下 overhead 高，默认精度设置易导致结果偏差

加速方案建议 #

1. Numpy运算优化 #

• 彻底消除Python循环：所有逐元素操作改用Numpy向量化函数，利用C级运算提升速度
• 预计算常量：提前计算coil_mat.T、1/(n_points*factor)等固定值，避免迭代时重复计算
• 内存布局优化：将coil_mat转换为连续内存布局coil_mat = np.ascontiguousarray(coil_mat)，提升矩阵乘法效率

2. Numba即时编译加速 #

用Numba将目标函数和Jacobian函数编译为机器码，大幅降低Python解释器开销：

from numba import jit

@jit(nopython=True)
def _residuals_mse_numba(coef, unshimmed_vec, coil_mat, factor, reg_factor, reg_factor_channel):
    residual = unshimmed_vec + coil_mat @ coef
    mse_term = np.mean(residual ** 2) / factor
    l1_term = reg_factor * np.mean(np.abs(coef) / reg_factor_channel)
    return mse_term + l1_term

@jit(nopython=True)
def _residuals_mse_jacobian_numba(coef, unshimmed_vec, coil_mat, factor, reg_factor, reg_factor_channel):
    n_points = unshimmed_vec.size
    b = 2.0 / (n_points * factor)
    residual = unshimmed_vec + coil_mat @ coef
    mse_grad = b * coil_mat.T @ residual
    l1_grad = reg_factor * np.sign(coef) / reg_factor_channel / coef.size
    return mse_grad + l1_grad

首次调用会有编译开销，后续重复调用速度提升显著

3. 适配问题的专用求解器 #

将问题转换为标准正则化线性回归形式，使用sklearn的高效求解器：

• 目标等价于最小化：(1/(2n*factor)) ||coil_mat @ coef + unshimmed_vec||² + (reg_factor/m) sum(|coef_j|/reg_factor_channel_j)
• 可使用sklearn.linear_model.SGDRegressor自定义损失函数，适合大规模数据的快速迭代；或用ElasticNet适配正则项逻辑

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