使用scipy minimize时出现超大数组内存分配错误求助

问题情况 #

需要优化维度为(300,300)的变量W，以最小化Frobenius范数||X - WX||_F，但第一次迭代后出现内存错误：无法为形状(90000,90000)的数组分配60.3 GiB内存，理论上该数组不应存在。尝试多种优化方法均未解决问题。

原可复现代码：

import numpy as np
from scipy.optimize import minimize

def objective_function(W, X, d):
    W = W.reshape((d, d)) 
    return np.linalg.norm(X - np.dot(W, X), 'fro')

def estimate_background(Y, A_hat, C_hat, l, d):
    A_hat = A_hat.T
    T = Y.shape[1] # T == 120
    b0 = (Y - A_hat @ C_hat) @ np.ones((T, 1))
    X = Y - np.dot(A_hat, C_hat) - np.dot(b0, np.ones((1, T))) # X.shape == (300,120)
    W_init = np.zeros((300,300)) # d == 300
    result = minimize(objective_function, x0=W_init, args=(X, d))

    optimal_W = result.x.reshape((d, d))

    return optimal_W

Y_tmp = np.ones((300,120))
A_tmp = np.ones((3,300))
C_tmp = np.ones((3,120))
B = estimate_background(Y_tmp, A_tmp, C_tmp, 5, Y_tmp.shape[0])

错误信息：

MemoryError: Unable to allocate 60.3 GiB for an array with shape (90000, 90000) and data type float64

原因分析 #

scipy的minimize默认使用L-BFGS-B优化方法，当未手动提供梯度（jac参数为None）时，会采用数值梯度计算。此时，输入的W被扁平化为90000个元素的一维数组，数值梯度需要为每个元素计算一次函数值，进而构建一个(90000,90000)的Jacobian矩阵，这直接导致内存溢出。

解决方案 #

核心是避免数值梯度的计算，手动推导并提供解析梯度，大幅降低内存开销。

梯度推导 #

目标函数等价于最小化||X - WX||_F^2（范数平方的最小值点与原范数一致，计算更简便）。展开后对W求导可得：
∇W = 2*(W @ X @ X.T - X @ X.T)
将该梯度扁平化后即可作为优化的梯度输入。

修改后的代码 #

import numpy as np
from scipy.optimize import minimize

def objective_function(W, X, d):
    W = W.reshape((d, d))
    residual = X - W @ X
    # 返回函数值和梯度（同时返回可提升效率）
    loss = np.linalg.norm(residual, 'fro')
    # 计算梯度并扁平化
    grad = 2 * (W @ X @ X.T - X @ X.T).flatten()
    return loss, grad

def estimate_background(Y, A_hat, C_hat, l, d):
    A_hat = A_hat.T
    T = Y.shape[1]
    b0 = (Y - A_hat @ C_hat) @ np.ones((T, 1))
    X = Y - A_hat @ C_hat - b0 @ np.ones((1, T))
    W_init = np.zeros((d, d)).flatten()  # 直接传入扁平化的初始值
    # 指定jac=True，表示目标函数返回(损失, 梯度)
    result = minimize(objective_function, x0=W_init, args=(X, d), jac=True)

    optimal_W = result.x.reshape((d, d))
    return optimal_W

# 测试
Y_tmp = np.ones((300,120))
A_tmp = np.ones((3,300))
C_tmp = np.ones((3,120))
B = estimate_background(Y_tmp, A_tmp, C_tmp, 5, Y_tmp.shape[0])

说明 #

1. 手动计算梯度避免了数值梯度带来的巨大内存开销，梯度仅为90000个元素的一维数组，内存占用约720KB（float64类型）。
2. 设置jac=True告知minimize函数目标函数会同时返回损失值和梯度，无需自动计算数值梯度。
3. 初始值W_init直接扁平化传入，减少一次形状转换的开销。

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