如何优化所实现的矩阵分解代码？存在哪些性能瓶颈？

我正在实现一段矩阵分解的推荐系统代码，想请教如何进一步优化该代码，以及当前版本的性能瓶颈具体位于哪些位置？

原代码如下：

import numpy

def matrix_factorization(R, P, Q, K, steps=5000, alpha=0.0002, beta=0.02):
   
    Q = Q.T

    for step in range(steps):
        for i in range(len(R)):
            for j in range(len(R[i])):
                if R[i][j] > 0:
                    # calculate error
                    eij = R[i][j] - numpy.dot(P[i,:],Q[:,j])

                    for k in range(K):
                        # calculate gradient with a and beta parameter
                        P[i][k] = P[i][k] + alpha * (2 * eij * Q[k][j] - beta * P[i][k])
                        Q[k][j] = Q[k][j] + alpha * (2 * eij * P[i][k] - beta * Q[k][j])

        eR = numpy.dot(P,Q)

        e = 0

        for i in range(len(R)):

            for j in range(len(R[i])):

                if R[i][j] > 0:

                    e = e + pow(R[i][j] - numpy.dot(P[i,:],Q[:,j]), 2)

                    for k in range(K):

                        e = e + (beta/2) * (pow(P[i][k],2) + pow(Q[k][j],2))
        # 0.001: local minimum
        if e < 0.001:

            break

    return P, Q.T

性能瓶颈分析 #

• 多层嵌套Python循环：代码包含4层嵌套循环（迭代轮次→行索引i→列索引j→隐因子维度k），Python解释器处理循环的效率极低，当R的规模较大（如用户/物品数过万）时，这会成为致命性能瓶颈。
• 冗余计算：每轮迭代中，更新P/Q时已计算过numpy.dot(P[i,:], Q[:,j])，后续计算总误差e时又重复计算相同点积，造成不必要的开销。
• 全量遍历稀疏矩阵：推荐系统的评分矩阵R通常是稀疏的（大部分元素为0，代表未评分），但原代码遍历所有i和j，大量时间浪费在处理无意义的0值上。
• 低效的误差计算：总误差e的计算采用三层嵌套循环逐元素累加，远不如NumPy向量化运算高效，且每轮迭代都执行该计算进一步拖慢速度。

优化方案 #

1. 向量化运算替代嵌套循环 #

利用NumPy的矩阵运算能力，将Python循环替换为底层优化的C语言实现操作，大幅提升速度。核心优化逻辑示例：

import numpy

def matrix_factorization_vectorized(R, P, Q, K, steps=5000, alpha=0.0002, beta=0.02):
    # 提前提取所有非零评分的索引
    non_zero_rows, non_zero_cols = numpy.where(R > 0)
    
    for step in range(steps):
        # 计算所有非零位置的预测值与误差
        predictions = numpy.dot(P, Q.T)
        errors = R[non_zero_rows, non_zero_cols] - predictions[non_zero_rows, non_zero_cols]
        
        # 向量化更新P和Q
        P[non_zero_rows] += alpha * (2 * errors[:, numpy.newaxis] * Q[non_zero_cols] - beta * P[non_zero_rows])
        Q[non_zero_cols] += alpha * (2 * errors[:, numpy.newaxis] * P[non_zero_rows] - beta * Q[non_zero_cols])
        
        # 向量化计算总误差
        pred_errors = R[non_zero_rows, non_zero_cols] - numpy.dot(P, Q.T)[non_zero_rows, non_zero_cols]
        total_error = numpy.sum(pred_errors ** 2) + (beta/2) * (numpy.sum(P ** 2) + numpy.sum(Q ** 2))
        
        if total_error < 0.001:
            break
    
    return P, Q

2. 利用Numba进行JIT编译 #

若无法完全避免循环，可用Numba对函数进行即时编译，将Python代码转换为机器码执行，大幅提升循环效率：

import numpy
from numba import jit

@jit(nopython=True)
def matrix_factorization_numba(R, P, Q, K, steps=5000, alpha=0.0002, beta=0.02):
    Q = Q.T
    for step in range(steps):
        for i in range(R.shape[0]):
            for j in range(R.shape[1]):
                if R[i][j] > 0:
                    eij = R[i][j] - numpy.dot(P[i,:], Q[:,j])
                    for k in range(K):
                        P[i][k] += alpha * (2 * eij * Q[k][j] - beta * P[i][k])
                        Q[k][j] += alpha * (2 * eij * P[i][k] - beta * Q[k][j])
        # 向量化计算误差，替代嵌套循环
        pred = numpy.dot(P, Q)
        mask = R > 0
        e = numpy.sum((R[mask] - pred[mask]) ** 2) + (beta/2) * (numpy.sum(P**2) + numpy.sum(Q**2))
        if e < 0.001:
            break
    return P, Q.T

3. 优化迭代终止逻辑 #

• 不需要每轮都计算精确总误差，可改为每隔N轮（如10轮）计算一次，减少计算开销；
• 监控误差下降幅度，当单轮下降量小于阈值时提前终止，避免无效迭代。

4. 超大规模数据场景：GPU加速 #

若处理超大规模评分矩阵，可改用TensorFlow、PyTorch等框架实现矩阵分解，利用GPU的并行计算能力进一步提升速度。

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