python - 在 python、go 或 julia 中快速直接访问像素

Question

我编写了一个小程序，可以创建随机噪声并将其全屏显示(5K 分辨率)。我用 pygame 来实现。但刷新速度却非常慢。 surfarray.blit_array 和随机生成都需要花费大量时间。有什么办法可以加快这个速度吗？我也可以灵活地使用 julia 或 golang 来代替。或者还有psychopy或octave与psychotoolbox(但是这些似乎不能在linux/wayland下工作)。

这是我写的:

import pygame
import numpy as N
import pygame.surfarray as surfarray
from numpy import int32, uint8, uint

 
def main():
     
    pygame.init()
     
    #flags = pygame.OPENGL | pygame.FULLSCREEN   # OpenGL does not want to work with surfarray
    flags = pygame.FULLSCREEN
    screen = pygame.display.set_mode((0,0), flags=flags, vsync=1)
    w, h = screen.get_width(), screen.get_height()

    clock = pygame.time.Clock()
    font = pygame.font.SysFont("Arial" , 18 , bold = True)
     
    # define a variable to control the main loop
    running = True

    def fps_counter():
        fps = str(int(clock.get_fps()))
        fps_t = font.render(fps , 1, pygame.Color("RED"))
        screen.blit(fps_t,(0,0))

                    
     
    # main loop
    while running:
        # event handling, gets all event from the event queue
        for event in pygame.event.get():
            # only do something if the event is of type QUIT
            if event.type == pygame.QUIT:
                # change the value to False, to exit the main loop
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    pygame.quit()
                    return
        array_img = N.random.randint(0, high=100, size=(w,h,3), dtype=uint)
        surfarray.blit_array(screen, array_img)
        fps_counter()
        pygame.display.flip()
        clock.tick()
        #print(clock.get_fps())
     
# run the main function only if this module is executed as the main script
# (if you import this as a module then nothing is executed)
if __name__=="__main__":
    # call the main function
    main()

我需要至少 30 fps 的刷新率才能发挥作用

Answer 1

更快的随机数生成

生成随机数的成本很高。当随机数生成器 (RNG) 需要统计上准确(即随机数即使在经过一些变换后也需要看起来非常随机)以及按顺序生成数字时尤其如此。

事实上，对于加密用途或某些数学(蒙特卡罗)模拟，目标 RNG 需要足够先进，以便几个后续生成的随机数之间不存在统计相关性。实际上，执行此操作的软件方法非常昂贵，以至于现代主流处理器提供了一种执行此操作的方法 specific instructions 。但并非所有处理器都支持这一点，并且 AFAIK Numpy 不使用它(当然是为了可移植性，因为在多台机器上具有相同种子的随机序列预计会给出相同的结果)。

幸运的是，RNG 在大多数其他用例中通常不需要那么准确。他们只需要看起来很随意。有many methods来做到这一点(例如 Mersenne Twister、Xoshiro、Xorshift、PCG/LCG)。 RNG 的性能、准确性和特化之间通常需要权衡。由于 Numpy 需要提供相对准确的通用 RNG(尽管据我所知并不意味着用于加密用例)，因此其性能次优也就不足为奇了。

对许多不同方法进行了有趣的回顾 here (尽管结果应该持保留态度，特别是在性能方面，因为 SIMD 友好等细节对于许多用例中的性能至关重要)。

在纯 Python(使用 CPython)中实现非常快速的随机数生成器是不可能的，但可以使用 Numba (或 Cython)来做到这一点。不过，可能有一些用 native 语言编写的快速现有模块可以做到这一点。最重要的是，我们可以使用多个线程来加速操作。为了简单起见(也因为它相对较快)，我选择实现 Xorshift64 RNG。

import numba as nb

@nb.njit('uint64(uint64,)')
def xorshift64_step(seed):
    seed ^= seed << np.uint64(13)
    seed ^= seed >> np.uint64(7)
    seed ^= seed << np.uint64(17)
    return seed

@nb.njit('uint64()')
def init_xorshift64():
    seed = np.uint64(np.random.randint(0x10000000, 0x7FFFFFFF)) # Bootstrap
    return xorshift64_step(seed)

@nb.njit('(uint64, int_)')
def random_pixel(seed, high):
    # Must be a constant for sake of performance and in the range [0;256]
    max_range = np.uint64(high)
    # Generate 3 group of 16 bits from the RNG
    bits1 = seed & np.uint64(0xFFFF)
    bits2 = (seed >> np.uint64(16)) & np.uint64(0xFFFF)
    bits3 = seed >> np.uint64(48)
    # Scale the numbers using a trick to avoid a modulo 
    # (since modulo are inefficient and statistically incorrect here)
    r = np.uint8(np.uint64(bits1 * max_range) >> np.uint64(16))
    g = np.uint8(np.uint64(bits2 * max_range) >> np.uint64(16))
    b = np.uint8(np.uint64(bits3 * max_range) >> np.uint64(16))
    new_seed = xorshift64_step(seed)
    return (r, g, b, new_seed)

@nb.njit('(int_, int_, int_)', parallel=True)
def pseudo_random_image(w, h, high):
    res = np.empty((w, h, 3), dtype=np.uint8)
    for i in nb.prange(w):
        # RNG seed initialization
        seed = init_xorshift64()
        for j in range(h):
            r, g, b, seed = random_pixel(seed, high)
            res[i, j, 0] = r
            res[i, j, 1] = g
            res[i, j, 2] = b
    return res

代码相当大，但在我的 6 核 i5-9600KF CPU 上，它比 Numpy 快约 22 倍。请注意，类似的代码可以在 Julia 中使用，以便获得快速实现(因为 Julia 使用类似于 Numba 的基于 LLVM 的 JIT)。

在我的机器上，这足以达到75 FPS(最大值)，而初始代码达到 16 FPS。

---

更快的操作和渲染

在大多数平台上，生成新的随机数组受到页面错误速度的限制。这会显着减慢计算速度。在 Python 中缓解这种情况的唯一方法是创建一次 Brame 缓冲区并执行就地操作。此外，PyGame 当然会在内部执行复制(也可能执行许多绘制调用)，因此使用较低级别的 API 可以显着加快速度。尽管如此，该操作可能会受到内存限制，并且没有什么可以避免的。但此时对您来说可能已经足够快了。

最重要的是，帧在 GPU 上渲染，因此 CPU 需要发送/复制 GPU 上的缓冲区，通常通过独立 GPU 的 PCIe 互连。对于宽屏幕，此操作不是很快。

实际上，您可以使用着色器直接在 GPU 上生成随机图像(或 OpenCL/CUDA 等工具)。这避免了上述开销，并且 GPU 可以比 CPU 更快地完成此任务。