def __init__(self, path, batch_size, path_filter="*.npy", num_threads=1, device_id=0, num_gpus=1):
super(NumpyReaderPipeline, self).__init__(batch_size,
num_threads,
device_id)
self.input = ops.NumpyReader(file_root = path,
file_filter = path_filter,
shard_id = device_id,
num_shards = num_gpus)
def get_numpy_reader(files, shard_id, num_shards, seed, shuffle):
return ops.NumpyReader(
seed=seed,
files=files,
device="cpu",
read_ahead=True,
shard_id=shard_id,
pad_last_batch=True,
num_shards=num_shards,
dont_use_mmap=True,
shuffle_after_epoch=shuffle,
)
def get_numpy_reader(files, shard_id, num_shards, seed, shuffle):
if shard_id is None:
shard_id = int(os.getenv("LOCAL_RANK", "0"))
return ops.NumpyReader(
seed=seed,
files=files,
device="cpu",
read_ahead=True,
shard_id=shard_id,
pad_last_batch=True,
num_shards=num_shards,
dont_use_mmap=True,
shuffle_after_epoch=shuffle,
)