Function test_pad_last_sample

dali/test/python/reader/test_numpy.py:984–1031

(device, batch_description, dont_use_mmap, use_o_direct)

@cartesian_params(("cpu", "gpu"), ((1, 2, 1), (3, 1, 2)), (True, False), (True, False))
def test_pad_last_sample(device, batch_description, dont_use_mmap, use_o_direct):
    if not is_gds_supported() and device == "gpu":
        raise SkipTest("GDS is not supported in this platform")
    if not dont_use_mmap and use_o_direct:
        raise SkipTest("Cannot use O_DIRECT with mmap")
    with tempfile.TemporaryDirectory(prefix=gds_data_root) as test_data_root:
        # create files
        num_samples, batch_size, num_shards = batch_description
        filenames = []
        ref_filenames = []
        arr_np_list = []
        last_file_name = None
        for index in range(0, num_samples):
            filename = os.path.join(test_data_root, "test_{:02d}.npy".format(index))
            last_file_name = filename
            filenames.append(filename)
            create_numpy_file(filename, (5, 2, 8), np.float32, False)
            arr_np_list.append(np.load(filename))
            ref_filenames.append(filename)
        while len(arr_np_list) < batch_size:
            arr_np_list.append(np.load(last_file_name))
            ref_filenames.append(last_file_name)
        pipe = NumpyReaderPipeline(
            path=test_data_root,
            files=filenames,
            file_list=None,
            file_filter=None,
            device=device,
            batch_size=batch_size,
            num_threads=4,
            device_id=0,
            pad_last_batch=True,
            num_shards=num_shards,
            dont_use_mmap=dont_use_mmap,
            enable_o_direct=use_o_direct,
        )

        try:
            for _ in range(2):
                pipe_out = pipe.run()
                for i in range(batch_size):
                    out_arr = to_array(pipe_out[0][i])
                    out_prop = pipe_out[0][i].source_info()
                    ref_arr = arr_np_list[i]
                    assert out_prop == ref_filenames[i]
                    assert_array_equal(out_arr, ref_arr)
        finally:
            del pipe


@cartesian_params(("global", "local", "none"), (True, False))