Markov链模拟：Python+Numpy较Matlab慢一个数量级的瓶颈分析与提速方法

我在Python+Numpy和Matlab中运行逻辑完全一致的马尔可夫链测试代码，发现Matlab版本的运行速度比Python快一个数量级。希望能定位Python代码的性能瓶颈，并找到有效的提速方法。

测试代码 #

Python+Numpy测试代码 #

# Packages
import numpy as np
import time

# Number of possible outcomes
num_outcomes = 20
# Dimension of the system
dim = 50
# Number of iterations
num_iterations = int(1e7)

# Possible outcomes
outcomes = np.arange(num_outcomes)
# Possible transition matrices
matrices = [np.random.rand(dim, dim) for k in outcomes]
matrices = [mat/np.sum(mat, axis=0) for mat in matrices]
# Initial state
state = np.random.rand(dim)
state = state/np.sum(state)

# List of samples
samples = np.random.choice(outcomes, size=(num_iterations,))
samples = samples.tolist()


# === PERFORMANCE-SENSITIVE PART OF THE CODE ===

# Update the state over all iterations
start_time = time.time()
for k in range(num_iterations):
    sample = samples[k]
    matrix = matrices[sample]
    state = np.matmul(matrix, state)
end_time = time.time()

# Print the execution time
print(end_time - start_time)

Matlab测试代码 #

% Number of possible outcomes
num_outcomes = 20;
% Number of dimensions
dim = 50;
% Number of iterations
num_iterations = 1e7;

% Possible outcomes
outcomes = 1:num_outcomes;
% Possible transition matrices
matrices = rand(num_outcomes, dim, dim);
matrices = matrices./sum(matrices,2);
matrices = num2cell(matrices,[2,3]);
matrices = cellfun(@shiftdim, matrices, 'UniformOutput', false);
% Initial state
state = rand(dim,1);
state = state./sum(state);

% List of samples
samples = datasample(outcomes, num_iterations);


% === PERFORMANCE-SENSITIVE PART OF THE CODE ===

% Update the state over all iterations
tic;
for k = 1:num_iterations
    sample = samples(k);
    matrix = matrices{sample};
    state = matrix * state;
end
toc;

测试环境 #

• Python：3.10 + Numpy 1.22.4
• Matlab：R2022a
• 系统：Windows 11 Pro 64位
• 硬件：联想ThinkPad T14，处理器为11th Gen Intel(R) Core(TM) i7-1165G7 @ 2.80GHz（4核8线程）

后续测试结论 #

1. 编辑1测试结果：通过额外测试发现，性能差异的核心原因是低矩阵尺寸下Python存在特有的固定调用开销——随着矩阵维度增大，Python的相对性能会逐渐接近Matlab，但小维度场景下这种固定开销被放大，导致整体速度慢一个数量级。这意味着JIT编译器可能能缓解这类开销问题。

2. 编辑2测试结果：在PyPy 3.9-v7.3.11-win64环境下运行代码后发现，PyPy改变了性能缩放比例，甚至在小矩阵尺寸下性能优于CPython，但针对这段特定代码，整体运行开销仍然较高，没有达到Matlab的水平。

可行优化方向 #

• JIT编译优化：使用Numba对性能敏感的循环部分进行JIT编译，消除Python循环的固定开销，提升小矩阵维度下的执行效率。
• Cython实现：将核心循环逻辑用Cython重写，通过静态类型声明和编译为机器码，大幅降低Python层面的调用开销。
• 批量运算重构：尝试重构代码逻辑，将多次小矩阵-向量乘法合并为批量运算，减少循环次数，利用Numpy的向量化优势抵消单步调用开销。

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