def test_pax_format(): global test_batch_size num_samples = 1000 tar_file_path = os.path.join(get_dali_extra_path(), "db/webdataset/MNIST/devel-0.tar") pax_tar_file_path = os.path.join(get_dali_extra_path(), "db/webdataset/pax/devel-0.tar") index_file = generate_temp_index_file(tar_file_path) num_shards = 100 for shard_id in range(num_shards): compare_pipelines( webdataset_raw_pipeline( tar_file_path, index_file.name, ["jpg", "cls"], num_shards=num_shards, shard_id=shard_id, batch_size=test_batch_size, device_id=0, num_threads=1, ), webdataset_raw_pipeline( pax_tar_file_path, None, ext=["jpg", "cls"], num_shards=num_shards, shard_id=shard_id, batch_size=test_batch_size, device_id=0, num_threads=1, ), test_batch_size, math.ceil(num_samples / num_shards / test_batch_size) * 2, )
def general_corner_case(test_batch_size=base.test_batch_size, dtypes=None, missing_component_behavior="", **kwargs): num_samples = 1000 tar_file_path = os.path.join(get_dali_extra_path(), "db/webdataset/MNIST/devel-0.tar") index_file = base.generate_temp_index_file(tar_file_path) extract_dir = base.generate_temp_extract(tar_file_path) equivalent_files = sorted( glob(extract_dir.name + "/*"), key=lambda s: int(s[s.rfind("/") + 1:s.rfind(".")])) compare_pipelines( base.webdataset_raw_pipeline( tar_file_path, index_file.name, ["jpg", "cls"], missing_component_behavior=missing_component_behavior, dtypes=dtypes, batch_size=test_batch_size, device_id=0, num_threads=1, **kwargs), base.file_reader_pipeline(equivalent_files, ["jpg", "cls"], batch_size=test_batch_size, device_id=0, num_threads=1, **kwargs), test_batch_size, math.ceil(num_samples / test_batch_size), )
def test_return_empty(): global test_batch_size num_samples = 1000 tar_file_path = os.path.join(get_dali_extra_path(), "db/webdataset/MNIST/missing.tar") index_file = generate_temp_index_file(tar_file_path) extract_dir = generate_temp_extract(tar_file_path) equivalent_files = glob(extract_dir.name + "/*") equivalent_files = sorted( equivalent_files, key=(lambda s: int(s[s.rfind("/") + 1:s.rfind(".")]))) # noqa: 203 compare_pipelines( webdataset_raw_pipeline( tar_file_path, index_file.name, ["jpg", "txt"], batch_size=test_batch_size, device_id=0, num_threads=1, missing_component_behavior="empty", ), file_reader_pipeline(equivalent_files, ["jpg", []], batch_size=test_batch_size, device_id=0, num_threads=1), test_batch_size, math.ceil(num_samples / test_batch_size), )
def test_dtypes(): global test_batch_size num_samples = 100 tar_file_path = os.path.join(get_dali_extra_path(), "db/webdataset/sample-tar/dtypes.tar") index_file = generate_temp_index_file(tar_file_path) wds_pipeline = webdataset_raw_pipeline( tar_file_path, index_file.name, ["float16", "int32", "float64"], dtypes=[dali.types.FLOAT16, dali.types.INT32, dali.types.FLOAT64], batch_size=test_batch_size, device_id=0, num_threads=1, ) wds_pipeline.build() for sample_idx in range(num_samples): if sample_idx % test_batch_size == 0: f16, i32, f64 = wds_pipeline.run() assert (f16.as_array()[sample_idx % test_batch_size] == [float(sample_idx)] * 10).all() assert (i32.as_array()[sample_idx % test_batch_size] == [int(sample_idx)] * 10).all() assert (f64.as_array()[sample_idx % test_batch_size] == [float(sample_idx)] * 10).all()
def test_raise_error_on_missing(): global test_batch_size tar_file_path = os.path.join(get_dali_extra_path(), "db/webdataset/MNIST/missing.tar") index_file = generate_temp_index_file(tar_file_path) wds_pipeline = webdataset_raw_pipeline( tar_file_path, index_file.name, ["jpg", "cls"], missing_component_behavior="error", batch_size=test_batch_size, device_id=0, num_threads=1, ) assert_raises(RuntimeError, wds_pipeline.build, glob="Underful sample detected")
def test_skip_sample(): global test_batch_size num_samples = 500 tar_file_path = os.path.join(get_dali_extra_path(), "db/webdataset/MNIST/missing.tar") index_file = generate_temp_index_file(tar_file_path) extract_dir = generate_temp_extract(tar_file_path) equivalent_files = list( filter( lambda s: int(s[s.rfind("/") + 1:s.rfind(".")]) < 2500, # noqa: 203 sorted(glob(extract_dir.name + "/*"), key=lambda s: int(s[s.rfind("/") + 1:s.rfind(".")]) ), # noqa: 203 )) compare_pipelines( webdataset_raw_pipeline( tar_file_path, index_file.name, ["jpg", "cls"], missing_component_behavior="skip", batch_size=test_batch_size, device_id=0, num_threads=1, ), file_reader_pipeline(equivalent_files, ["jpg", "cls"], batch_size=test_batch_size, device_id=0, num_threads=1), test_batch_size, math.ceil(num_samples / test_batch_size), ) wds_pipeline = webdataset_raw_pipeline( tar_file_path, index_file.name, ["jpg", "cls"], missing_component_behavior="skip", batch_size=test_batch_size, device_id=0, num_threads=1, ) wds_pipeline.build() assert_equal(list(wds_pipeline.epoch_size().values())[0], num_samples)
def test_wds_sharding(): global test_batch_size num_samples = 3000 tar_file_paths = [ os.path.join(get_dali_extra_path(), "db/webdataset/MNIST/devel-0.tar"), os.path.join(get_dali_extra_path(), "db/webdataset/MNIST/devel-1.tar"), os.path.join(get_dali_extra_path(), "db/webdataset/MNIST/devel-2.tar"), ] index_files = [ generate_temp_index_file(tar_file_path) for tar_file_path in tar_file_paths ] extract_dirs = [ generate_temp_extract(tar_file_path) for tar_file_path in tar_file_paths ] equivalent_files = sum( list( sorted(glob(extract_dir.name + "/*"), key=lambda s: int(s[s.rfind("/") + 1:s.rfind(".")]) ) # noqa: 203 for extract_dir in extract_dirs), [], ) compare_pipelines( webdataset_raw_pipeline( tar_file_paths, [index_file.name for index_file in index_files], ["jpg", "cls"], batch_size=test_batch_size, device_id=0, num_threads=1, ), file_reader_pipeline( equivalent_files, ["jpg", "cls"], batch_size=test_batch_size, device_id=0, num_threads=1, ), test_batch_size, math.ceil(num_samples / test_batch_size), )
def test_wide_sample(): test_batch_size = 1 num_samples = 1 tar_file_path = os.path.join(get_dali_extra_path(), "db/webdataset/sample-tar/wide.tar") index_file = base.generate_temp_index_file(tar_file_path) extract_dir = base.generate_temp_extract(tar_file_path) equivalent_files = list(sorted(glob(extract_dir.name + "/*"))) num_components = 1000 compare_pipelines( base.webdataset_raw_pipeline( tar_file_path, index_file.name, [str(x) for x in range(num_components)], batch_size=test_batch_size, device_id=0, num_threads=1, ), base.file_reader_pipeline( equivalent_files, [str(x) for x in range(num_components)], batch_size=test_batch_size, device_id=0, num_threads=1, ), test_batch_size, math.ceil(num_samples / test_batch_size) * 10, ) wds_pipeline = base.webdataset_raw_pipeline( tar_file_path, index_file.name, ["txt"], batch_size=test_batch_size, device_id=0, num_threads=1, ) wds_pipeline.build() assert_equal(list(wds_pipeline.epoch_size().values())[0], num_samples)
def test_sharding(): global test_batch_size num_samples = 1000 tar_file_path = os.path.join(get_dali_extra_path(), "db/webdataset/MNIST/devel-0.tar") index_file = generate_temp_index_file(tar_file_path) extract_dir = generate_temp_extract(tar_file_path) equivalent_files = sorted( glob(extract_dir.name + "/*"), key=lambda s: int(s[s.rfind("/") + 1:s.rfind(".")])) # noqa: 203 num_shards = 100 for shard_id in range(num_shards): compare_pipelines( webdataset_raw_pipeline( tar_file_path, index_file.name, ["jpg", "cls"], num_shards=num_shards, shard_id=shard_id, batch_size=test_batch_size, device_id=0, num_threads=1, ), file_reader_pipeline( equivalent_files, ["jpg", "cls"], num_shards=num_shards, shard_id=shard_id, batch_size=test_batch_size, device_id=0, num_threads=1, ), test_batch_size, math.ceil(num_samples / num_shards / test_batch_size) * 2, )