def pipe(max_batch_size, input_data, device): pipe = Pipeline(batch_size=max_batch_size, num_threads=4, device_id=0) data = fn.external_source(source=input_data, cycle=False, device=device, layout="FHWC") processed, _ = fn.element_extract(data, element_map=[0, 3]) pipe.set_outputs(processed) return pipe
def element_extract_pipe(shape, layout, element_map, dev, dtype): min_shape = [s // 2 if s > 1 else 1 for s in shape] min_shape[0] = shape[0] min_shape = tuple(min_shape) input = fn.external_source(source=RandomlyShapedDataIterator( batch_size, min_shape=min_shape, max_shape=shape, dtype=dtype), layout=layout) if dev == "gpu": input = input.gpu() elements = fn.element_extract(input, element_map=element_map) result = (input, ) + tuple(elements) if len(element_map) > 1 else ( input, elements) return result
def test_element_extract_cpu(): pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=None) test_data_shape = [5, 10, 20, 3] def get_data(): out = [ np.random.randint(0, 255, size=test_data_shape, dtype=np.uint8) for _ in range(batch_size) ] return out data = fn.external_source(source=get_data, layout="FHWC") processed, _ = fn.element_extract(data, element_map=[0, 3]) pipe.set_outputs(processed) pipe.build() for _ in range(3): pipe.run()
def dali_dataloader( tfrec_filenames, tfrec_idx_filenames, shard_id=0, num_shards=1, batch_size=128, num_threads=os.cpu_count(), image_size=224, num_workers=1, training=True): pipe = Pipeline(batch_size=batch_size, num_threads=num_threads, device_id=0) with pipe: inputs = fn.readers.tfrecord( path=tfrec_filenames, index_path=tfrec_idx_filenames, random_shuffle=training, shard_id=shard_id, num_shards=num_shards, initial_fill=10000, read_ahead=True, pad_last_batch=True, prefetch_queue_depth=num_workers, name='Reader', features={ 'image/encoded': tfrec.FixedLenFeature((), tfrec.string, ""), 'image/class/label': tfrec.FixedLenFeature([1], tfrec.int64, -1), }) jpegs = inputs["image/encoded"] if training: images = fn.decoders.image_random_crop( jpegs, device="mixed", output_type=types.RGB, random_aspect_ratio=[0.8, 1.25], random_area=[0.1, 1.0], num_attempts=100) images = fn.resize(images, device='gpu', resize_x=image_size, resize_y=image_size, interp_type=types.INTERP_TRIANGULAR) mirror = fn.random.coin_flip(probability=0.5) else: images = fn.decoders.image(jpegs, device='mixed', output_type=types.RGB) images = fn.resize(images, device='gpu', size=int(image_size / 0.875), mode="not_smaller", interp_type=types.INTERP_TRIANGULAR) mirror = False images = fn.crop_mirror_normalize( images.gpu(), dtype=types.FLOAT, crop=(image_size, image_size), mean=[0.485 * 255, 0.456 * 255, 0.406 * 255], std=[0.229 * 255, 0.224 * 255, 0.225 * 255], mirror=mirror) label = inputs["image/class/label"] - 1 # 0-999 label = fn.element_extract(label, element_map=0) # Flatten label = label.gpu() pipe.set_outputs(images, label) pipe.build() last_batch_policy = LastBatchPolicy.DROP if training else LastBatchPolicy.PARTIAL loader = DALIClassificationIterator( pipe, reader_name="Reader", auto_reset=True, last_batch_policy=last_batch_policy) return loader