Function benchmark_array_xd

benchmarks/benchmark_array_xd.py:92–138

()

def benchmark_array_xd():
    times = {}
    read_functions = (
        read_unformated,
        read_formatted_as_numpy,
        read_batch_unformated,
        read_batch_formatted_as_numpy,
        read_col_unformated,
        read_col_formatted_as_numpy,
    )
    with tempfile.TemporaryDirectory() as tmp_dir:
        feats = datasets.Features({"image": Array2D(SPEED_TEST_SHAPE, dtype="float32")})
        data = generate_examples(features=feats, num_examples=SPEED_TEST_N_EXAMPLES)
        times["write_array2d"] = write(feats, data, tmp_dir)
        for read_func in read_functions:
            times[read_func.__name__ + " after write_array2d"] = read_func(feats, tmp_dir)

    with tempfile.TemporaryDirectory() as tmp_dir:
        # don't use fixed length for fair comparison
        # feats = datasets.Features(
        #     {"image": datasets.Sequence(datasets.Sequence(datasets.Value("float32"), SPEED_TEST_SHAPE[1]), SPEED_TEST_SHAPE[0])}
        # )
        feats = datasets.Features({"image": datasets.Sequence(datasets.Sequence(datasets.Value("float32")))})
        data = generate_examples(
            features=feats, num_examples=SPEED_TEST_N_EXAMPLES, seq_shapes={"image": SPEED_TEST_SHAPE}
        )
        times["write_nested_sequence"] = write(feats, data, tmp_dir)
        for read_func in read_functions:
            times[read_func.__name__ + " after write_nested_sequence"] = read_func(feats, tmp_dir)

    with tempfile.TemporaryDirectory() as tmp_dir:
        # don't use fixed length for fair comparison
        # feats = datasets.Features(
        #     {"image": datasets.Sequence(datasets.Value("float32"), SPEED_TEST_SHAPE[0] * SPEED_TEST_SHAPE[1])}
        # )
        feats = datasets.Features({"image": datasets.Sequence(datasets.Value("float32"))})
        data = generate_examples(
            features=feats,
            num_examples=SPEED_TEST_N_EXAMPLES,
            seq_shapes={"image": [SPEED_TEST_SHAPE[0] * SPEED_TEST_SHAPE[1]]},
        )
        times["write_flattened_sequence"] = write(feats, data, tmp_dir)
        for read_func in read_functions:
            times[read_func.__name__ + " after write_flattened_sequence"] = read_func(feats, tmp_dir)

    with open(RESULTS_FILE_PATH, "wb") as f:
        f.write(json.dumps(times).encode("utf-8"))


if __name__ == "__main__":  # useful to run the profiler