马尔可夫链的模拟比在 Matlab 中慢

Question

我在 Python+Numpy 和 Matlab 中运行相同的测试代码，发现 Matlab 代码的速度快了一个数量级。我想知道 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 代码始终比 Matlab 代码慢一个数量级，我不确定为什么。

知道从哪里开始吗？

我使用 Python 3.10 解释器和 Numpy 1.22.4 运行 Python 代码。我使用 Matlab R2022a 运行 Matlab 代码。这两个代码都在具有以下处理器的 Lenovo T14 ThinkPad 上的 Windows 11 Pro 64 位上运行：

第 11 代 Intel(R) Core(TM) i7-1165G7 @ 2.80GHz，2803 Mhz，4 核，8 逻辑处理器

编辑 1：我做了一些额外的测试，看起来罪魁祸首是某种类型的 Python 特定的低矩阵大小的常量开销：

正如 hpaulj 和 MSS 所建议的，这可能意味着 JIT 编译器可以解决其中的一些问题。我会在不久的将来尽我所能尝试这个。

编辑 2：我在 Pypy 下运行了代码，虽然它确实改变了缩放比例，甚至在小矩阵尺寸下击败了 Cpython，但它通常会为这个特定代码带来很大的开销：

因此，如果有减轻这种开销的方法，JIT 编译器可能会有所帮助。否则，Cython 实现可能是剩下的路要走......

Answer 1

在循环中，主要障碍是np.matmul(matrix,state)。

如果我们展开循环：

st[1] = m[0]@st[0]
st[2] = m[1]@st[1] = m[1]@m[0]@st[0]
st[3] = m[2]@m[1]@m[0]@st[0]

没有明显的矢量化方式以非循环方式执行循环np.matmul。

更好的方法是在 log_2(n) 循环中执行此操作。

import numpy as np
outcomes = 20
dim = 50
num_iter = int(1e7)

mat = np.random.rand(outcomes,dim, dim)
mat = mat/mat.sum(axis=1)[...,None]
state = np.random.rand(dim)
state = state/np.sum(state)

samples = np.random.choice(np.arange(outcomes), size=(num_iter,))


a = mat[samples,...]
# This while loop takes log_2(num_iter) iterations
while len(a) > 1:
    a = np.matmul(a[::2, ...], a[1::2, ...])

state = np.matmul(a,state)

使用 numba jit 可以进一步减少时间。