def test_caffe2_reader_cpu():
pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=None)
out, _ = fn.caffe2_reader(path=caffe2_dir, shard_id=0, num_shards=1)
pipe.set_outputs(out)
pipe.build()
for _ in range(3):
pipe.run()
def mnist_pipeline(num_threads,
path,
device,
device_id=0,
shard_id=0,
num_shards=1,
seed=0):
pipeline = Pipeline(BATCH_SIZE, num_threads, device_id, seed)
with pipeline:
jpegs, labels = fn.caffe2_reader(path=path,
random_shuffle=True,
shard_id=shard_id,
num_shards=num_shards)
images = fn.image_decoder(jpegs,
device='mixed' if device == 'gpu' else 'cpu',
output_type=types.GRAY)
if device == 'gpu':
labels = labels.gpu()
images = fn.crop_mirror_normalize(images,
dtype=types.FLOAT,
mean=[0.],
std=[255.],
output_layout="CHW")
pipeline.set_outputs(images, labels)
return pipeline