Python for 循环 numpy 向量化

Question

您对如何加快下面代码中的“f_img_update”功能有什么建议吗？我尝试添加 numba @jit 装饰器，但没有运气。我想 numpy 矢量化可能会起作用，但我受的教育还不够，无法让它自己工作：/

代码：

import numpy as np
import time
import random
from random import randrange


dimensions = (100, 50)  # [m]
resolution = 5  # [cm] 1,2,4,5,10,20,25,50,100
pix_res = int(100 / resolution)

height = dimensions[0] * pix_res
width = dimensions[1] * pix_res
img_size = (height, width)

# Make "RGBA" image
img = np.zeros((*img_size, 4), np.uint16)

# DATA PREPARATION
xmin = 1700000
ymin = 1400000
zmin = 100

xmax = xmin + dimensions[0]
ymax = ymin + dimensions[1]
zmax = 150

data = []
for i in range(100000):
    xyzi = [random.uniform(xmin, xmax), random.uniform(ymin, ymax), random.uniform(zmin, zmax), randrange(255)]
    data.append(xyzi)

# IMAGE UPDATE
def f_img_update(f_img, f_xmin, f_ymin, f_data, f_pix_per_1m):  # IMAGE UPDATE
    for f_xyzi in f_data:
        f_x_img, f_y_img = f_xyzi[0:2]
        f_x_img = int((f_x_img - f_xmin) * 100 / f_pix_per_1m)
        f_y_img = int((f_y_img - f_ymin) * 100 / f_pix_per_1m)
        f_img[f_x_img][f_y_img] = f_xyzi
    return f_img


t1 = time.time()
img = f_img_update(img, xmin, ymin, data, pix_res)
t2 = time.time()
print("image update time: ", t2 - t1)

Answer 1

最简单的方法是将

data

从标准 python 列表转换为

np.array

，然后使用

@njit

装饰器：

import random
import time
from random import randrange

import numpy as np
from numba import njit

dimensions = (100, 50)  # [m]
resolution = 5  # [cm] 1,2,4,5,10,20,25,50,100
pix_res = int(100 / resolution)

height = dimensions[0] * pix_res
width = dimensions[1] * pix_res
img_size = (height, width)

# Make "RGBA" image
img = np.zeros((*img_size, 4), np.uint16)

# DATA PREPARATION
xmin = 1700000
ymin = 1400000
zmin = 100

xmax = xmin + dimensions[0]
ymax = ymin + dimensions[1]
zmax = 150

data = []
for i in range(100_000):
    xyzi = [
        random.uniform(xmin, xmax),
        random.uniform(ymin, ymax),
        random.uniform(zmin, zmax),
        randrange(255),
    ]
    data.append(xyzi)

data = np.array(data)


# IMAGE UPDATE
def f_img_update(f_img, f_xmin, f_ymin, f_data, f_pix_per_1m):  # IMAGE UPDATE
    for f_xyzi in f_data:
        f_x_img = f_xyzi[0]
        f_y_img = f_xyzi[1]
        f_x_img = int((f_x_img - f_xmin) * 100 / f_pix_per_1m)
        f_y_img = int((f_y_img - f_ymin) * 100 / f_pix_per_1m)
        f_img[f_x_img, f_y_img] = f_xyzi
    return f_img


@njit
def f_img_update_numba(f_img, f_xmin, f_ymin, f_data, f_pix_per_1m):  # IMAGE UPDATE
    for f_xyzi in f_data:
        f_x_img = f_xyzi[0]
        f_y_img = f_xyzi[1]
        f_x_img = int((f_x_img - f_xmin) * 100 / f_pix_per_1m)
        f_y_img = int((f_y_img - f_ymin) * 100 / f_pix_per_1m)
        f_img[f_x_img, f_y_img] = f_xyzi
    return f_img


# compile it first
assert np.allclose(
    f_img_update(img, xmin, ymin, data, pix_res),
    f_img_update_numba(img, xmin, ymin, data, pix_res),
)

img_copy = img.copy()
t1 = time.time()
_ = f_img_update(img_copy, xmin, ymin, data, pix_res)
t2 = time.time()
print("image update normal: ", t2 - t1)

img_copy = img.copy()
t1 = time.time()
_ = f_img_update_numba(img_copy, xmin, ymin, data, pix_res)
t2 = time.time()
print("image update time numba: ", t2 - t1)

在我的 AMD 5700x 上打印：

image update normal:  0.07558226585388184
image update time numba:  0.001255035400390625