首页 > 解决方案 > 为巨大的矩阵提高 numpy 嵌套 for 循环的速度

问题描述

我有一个 python 代码,它嵌套了 for 循环,用于一个巨大的矩阵。这里 rho0, alp0, bet0 是常数。E 是一个对称矩阵。我怎样才能提高for循环的速度?

N = 20000
Q = np.zeros(shape=(N,N))               
for i in tqdm(range(0,N)):
    for j in range(0,N):
        nom    = rho0*alp0**E[i,j]*(1-alp0)**(M-E[i,j]);
        dnom   = nom + ( (1-rho0)*bet0**E[i,j]*(1-bet0)**(M-E[i,j]));
        Q[i,j] = nom/dnom;

标签: python-3.xnumpy

解决方案


您可以通过将 Python 循环替换为硬件加速的向量运算以及由numpy.

(你没有定义几个变量,所以我冒昧地定义了它们。)

import numpy as np
from tqdm import tqdm

# Set these to the correct values.
rho = 0.6
alp0 = 1.4
bet0 = 2.5
M = 3.6
_E = np.random.rand(N, N)
E = _E + _E.T  # symmetric matrix.

# Look, all 20000 * 20000 values computed in one move!
nom = rho0 * alp0 ** E * (1 - alp0) ** (M - E)  # Shape is 20000x20000
dnim = nom + ((1 - rho0) * bet0 ** E * (1 - bet0) ** (M - E))  # Shape is 20000x20000
Q = nom / dnom

推荐阅读