如何提升Python中高维度矩阵乘法的执行效率？

Great question! Your pure Python nested loop implementation works fine for small matrices, but it'll get painfully slow with high-dimensional data—this is because Python's interpreted nature makes loops have significant overhead. Below are the most effective ways to speed up matrix multiplication for large matrices, ordered by practicality:

1. Use NumPy (Most Recommended CPU Solution) #

NumPy is the standard for scientific computing in Python. Its matrix operations are implemented in optimized C code, and it leverages BLAS/LAPACK backends (like OpenBLAS or MKL) to take full advantage of CPU multi-core processing and vector instruction sets.

Code Example #

import numpy as np
import time

# Input matrix dimensions
m = int(input("Enter rows of matrix A: "))
n = int(input("Enter columns of matrix A: "))
p = int(input("Enter rows of matrix B: "))
q = int(input("Enter columns of matrix B: "))

if n != p:
    print("Invalid entry: columns of A must equal rows of B")
    exit()

# Input matrix elements (or generate randomly for large matrices)
print("Enter matrix A elements (row-wise, space-separated):")
A = np.array([int(x) for x in input().split()]).reshape(m, n)
print("Enter matrix B elements (row-wise, space-separated):")
B = np.array([int(x) for x in input().split()]).reshape(p, q)

start_time = time.time()
# Perform matrix multiplication using the @ operator (or np.dot)
R = A @ B
print(R)
print("--- %s seconds ---" % (time.time() - start_time))

For large matrices (e.g., 1000x1000), NumPy will be hundreds to thousands of times faster than pure Python. If you use an MKL-optimized NumPy build (like the one included with Anaconda), you'll get an extra speed boost.

2. GPU Acceleration (Best for Ultra-Large Matrices) #

If you're working with extremely large matrices (e.g., 10000x10000), CPU performance might still not be enough. GPU acceleration can provide orders of magnitude faster computation:

• CuPy: Has an API nearly identical to NumPy, but runs operations on NVIDIA GPUs. It's perfect if you want a drop-in replacement for NumPy.
• PyTorch/TensorFlow: Deep learning frameworks with highly optimized matrix multiplication functions (torch.matmul/tf.matmul). Ideal if you need to follow up with other ML/DL operations.

CuPy Example #

import cupy as cp
import time

# Assume you already have NumPy arrays A_np and B_np
A_cp = cp.array(A_np)
B_cp = cp.array(B_np)

start_time = time.time()
R_cp = A_cp @ B_cp
cp.cuda.Stream.null.synchronize()  # Wait for GPU computation to finish
print("--- %s seconds ---" % (time.time() - start_time))
R_np = cp.asnumpy(R_cp)  # Convert back to CPU array if needed

3. Pure Python Optimization (Learning Only, Not Recommended for Production) #

If you absolutely must stick to pure Python, you can make small improvements, but the speed gain will be negligible compared to using libraries:

• Fix duplicate loop variables (your original code reuses i for multiple loops, which causes bugs)
• Adjust loop order to take advantage of CPU cache locality
• Use list comprehensions to reduce loop overhead

Corrected & Optimized Pure Python Code #

import time

m = int(input("Enter rows of A: "))
n = int(input("Enter columns of A: "))
p = int(input("Enter rows of B: "))
q = int(input("Enter columns of B: "))

if n != p:
    print("Invalid entry")
    exit()

# Input matrix A
print("Enter matrix A (row-wise, space-separated):")
A = [list(map(int, input().split())) for _ in range(m)]

# Input matrix B
print("Enter matrix B (row-wise, space-separated):")
B = [list(map(int, input().split())) for _ in range(p)]

start_time = time.time()
# Initialize result matrix with list comprehension
R = [[0 for _ in range(q)] for _ in range(m)]

# Reorder loops for better cache performance
for i in range(m):
    a_row = A[i]
    for k in range(n):
        a_ik = a_row[k]
        b_row = B[k]
        for j in range(q):
            R[i][j] += a_ik * b_row[j]

print(R)
print("--- %s seconds ---" % (time.time() - start_time))

Note: Your original code had a bug in initializing the result matrix R (reusing i for both row and column loops), which would produce an incorrectly shaped matrix. The code above fixes this issue.

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