import pytest import numpy as np import timeit from ding.data.shm_buffer import ShmBuffer import multiprocessing as mp def subprocess(shm_buf): data = np.random.rand(1024, 1024).astype(np.float32) res = timeit.repeat(lambda: shm_buf.fill(data), repeat=5, number=1000) print("Mean: {:.4f}s, STD: {:.4f}s, Mean each call: {:.4f}ms".format(np.mean(res), np.std(res), np.mean(res))) @pytest.mark.benchmark def test_shm_buffer(): data = np.random.rand(1024, 1024).astype(np.float32) shm_buf = ShmBuffer(data.dtype, data.shape, copy_on_get=False) proc = mp.Process(target=subprocess, args=[shm_buf]) proc.start() proc.join()