Блог

Bacground: Я храню коллекцию массивов numpy на твердотельном накопителе (около 500 тыс. файлов, достаточно больших, чтобы вместить все в память), и мне нужно загружать образцы пакетами. Поэтому, чтобы выяснить, какой формат является наиболее эффективным, я провел этот тест:

 import numpy as np
import pickle
from datetime import datetime
import random


def numpy_read_speed_test():

    root_folder = "C:\temp\"

    def pickle_read(pypic):
        with open(pypic, "rb") as pypif:
            return pickle.load(pypif)

    num_files = 1000

    file_names = [(root_folder   f"np_comp_{n}.npz",
                   root_folder   f"np_pickle_{n}.npy",
                   root_folder   f"np_flat_{n}.npy",
                   root_folder   f"picle_{n}.pkl") for n in range(num_files)]

    print(f"generating {num_files} test files")
    for np_comp, np_pickle, np_flat, py_pickle in file_names:

        # generate realistic dummy data
        arr = np.random.normal(0.0, 10.0, (144*16, 37))
        msk = np.random.randint(-1, 2, (1, 37))
        arr = (arr*msk).astype('float32')

        np.save(np_flat, arr, allow_pickle=False)
        np.save(np_pickle, arr, allow_pickle=True)
        np.savez_compressed(np_comp, x=arr)

        with open(py_pickle, "wb") as f:
            pickle.dump(arr, f)

    for _ in range(5):

        sample_size = num_files//2

        print(f"      testing numpy flat format: ", end='', flush=True)
        t0 = datetime.utcnow()
        npflt = [np.load(np_flat) for _, _, np_flat, _ in random.sample(file_names, sample_size)]
        t1 = datetime.utcnow()
        print(f"{t1 - t0} - {(t1 - t0).total_seconds()/sample_size} seconds per file", flush=True)

        print(f"testing numpy compressed format: ", end='', flush=True)
        t0 = datetime.utcnow()
        npcpz = [np.load(np_comp)['x'] for np_comp, _, _, _ in random.sample(file_names, sample_size)]
        t1 = datetime.utcnow()
        print(f"{t1 - t0} - {(t1 - t0).total_seconds()/sample_size} seconds per file", flush=True)

        print(f"   testing python pickle format: ", end='', flush=True)
        t0 = datetime.utcnow()
        pypkl = [pickle_read(py_pickle) for _, _, _, py_pickle in random.sample(file_names, sample_size)]
        t1 = datetime.utcnow()
        print(f"{t1 - t0} - {(t1 - t0).total_seconds()/sample_size} seconds per file", flush=True)

        print(f"testing numpy picle-able format: ", end='', flush=True)
        t0 = datetime.utcnow()
        npplk = [np.load(np_pickle) for _, np_pickle, _, _ in random.sample(file_names, sample_size)]
        t1 = datetime.utcnow()
        print(f"{t1 - t0} - {(t1 - t0).total_seconds()/sample_size} seconds per file", flush=True)

        print(f"lengths: {len(npplk)}, {len(npflt)}, {len(npcpz)}, {len(pypkl)}", flush=True)
        print(f"shapes: {npplk[0].shape}, {npflt[0].shape}, {npcpz[0].shape}, {pypkl[0].shape}", flush=True)
        print("", flush=True)


if __name__ == '__main__':
    numpy_read_speed_test()

 

Проблема в том, что тест дает противоречивые результаты:

 generating 1000 test files
      testing numpy flat format: 0:00:02.893651 - 0.005787302 seconds per file
testing numpy compressed format: 0:00:03.028680 - 0.00605736 seconds per file
   testing python pickle format: 0:00:01.840414 - 0.003680828 seconds per file
testing numpy picle-able format: 0:00:02.779625 - 0.00555925 seconds per file
lengths: 500, 500, 500, 500
shapes: (2304, 37), (2304, 37), (2304, 37), (2304, 37)

      testing numpy flat format: 0:00:01.682378 - 0.003364756 seconds per file
testing numpy compressed format: 0:00:02.062464 - 0.004124928 seconds per file
   testing python pickle format: 0:00:00.985221 - 0.001970442 seconds per file
testing numpy picle-able format: 0:00:01.518342 - 0.003036684 seconds per file
lengths: 500, 500, 500, 500
shapes: (2304, 37), (2304, 37), (2304, 37), (2304, 37)

      testing numpy flat format: 0:00:00.891200 - 0.0017824 seconds per file
testing numpy compressed format: 0:00:01.429321 - 0.002858642 seconds per file
   testing python pickle format: 0:00:00.515116 - 0.001030232 seconds per file
testing numpy picle-able format: 0:00:00.888200 - 0.0017764 seconds per file
lengths: 500, 500, 500, 500
shapes: (2304, 37), (2304, 37), (2304, 37), (2304, 37)

      testing numpy flat format: 0:00:00.448100 - 0.0008962 seconds per file
testing numpy compressed format: 0:00:01.226276 - 0.0024525519999999998 seconds per file
   testing python pickle format: 0:00:00.351079 - 0.0007021579999999999 seconds per file
testing numpy picle-able format: 0:00:00.699157 - 0.0013983140000000001 seconds per file
lengths: 500, 500, 500, 500
shapes: (2304, 37), (2304, 37), (2304, 37), (2304, 37)

      testing numpy flat format: 0:00:00.460103 - 0.000920206 seconds per file
testing numpy compressed format: 0:00:01.140257 - 0.002280514 seconds per file
   testing python pickle format: 0:00:00.292065 - 0.0005841300000000001 seconds per file
testing numpy picle-able format: 0:00:00.491111 - 0.000982222 seconds per file
lengths: 500, 500, 500, 500
shapes: (2304, 37), (2304, 37), (2304, 37), (2304, 37)
 

Существует большая разница в первом раунде и в следующем. Я полагаю, что это связано с системным кэшированием данных. Как бы вы улучшили тест?

Правка: Изначально была ошибка, из-за которой казалось, что вариация была очень высокой, но все равно похоже, что тест можно улучшить.