有沒有辦法加快以下代碼行的速度：
desired_channel=32
len_indices=50000
fast_idx = np.broadcast_to(np.arange(desired_channel)[:, None], (desired_channel, len_indices)).T.reshape(-1)
謝謝你。
uj5u.com熱心網友回復：
我是jax庫的新手。我使用Colab TPU上的以下代碼通過 jax one 比較了您的代碼：
import numpy as np
from jax import jit
import jax.numpy as jnp
import timeit
desired_channel=32
len_indices=50000
def ex_():
return np.broadcast_to(np.arange(desired_channel)[:, None], (desired_channel, len_indices)).T.reshape(-1)
%timeit -n1000 -r10 ex_()
@jit
def exj_():
return jnp.broadcast_to(jnp.arange(desired_channel)[:, None], (desired_channel, len_indices)).T.reshape(-1)
%timeit -n1000 -r10 exj_()
在我的一項努力中，結果如下：
1000 個回圈，最好的 10：每個回圈 901 μs
1000 個回圈，最好的 10：每個回圈 317 μs
通過這種方式，jax可以將您的代碼速度提高兩到三倍。
uj5u.com熱心網友回復：
最后一行代碼簡單地等于np.tile(np.arange(desired_channel), len_indices).
在我的機器上，np.tile許多 Numpy 呼叫的性能受作業系統（頁面錯誤）、記憶體分配器和記憶體吞吐量的限制。有兩種方法可以克服此限制：不分配/填充臨時緩沖區，使用較短的型別（例如或關于您的需要）在記憶體中生成較小的陣列。np.uint8np.uint16
由于函式沒有out引數np.tile，Numba 可用于生成快速替代函式。下面是一個例子：
import numba as nb
@nb.njit('int32[::1](int32, int32, int32[::1])', parallel=True)
def generate(desired_channel, len_indices, out):
for i in nb.prange(len_indices):
for j in range(desired_channel):
out[i*desired_channel j] = j
return out
desired_channel=32
len_indices=50000
buffer = np.full(desired_channel * len_indices, 0, dtype=np.int32)
%timeit -n 200 generate(desired_channel, len_indices, fast_idx)
以下是性能結果：
Original code: 1.25 ms
np.tile: 1.24 ms
Numba: 0.20 ms
轉載請註明出處，本文鏈接：https://www.uj5u.com/yidong/335483.html