def __init__(self, batch_size, num_threads, device_id):
super(TFRecordPipeline, self).__init__(batch_size, num_threads,
device_id)
self.input = ops.TFRecordReader(
path=tfrecord,
index_path=tfrecord_idx,
features={
"image/encoded": tfrec.FixedLenFeature((), tfrec.string, ""),
'image/class/label': tfrec.FixedLenFeature([1], tfrec.int64,
-1),
'image/class/text': tfrec.FixedLenFeature([], tfrec.string,
''),
'image/object/bbox/xmin':
tfrec.VarLenFeature(tfrec.float32, 0.0),
'image/object/bbox/ymin':
tfrec.VarLenFeature(tfrec.float32, 0.0),
'image/object/bbox/xmax':
tfrec.VarLenFeature(tfrec.float32, 0.0),
'image/object/bbox/ymax':
tfrec.VarLenFeature(tfrec.float32, 0.0)
})
self.decode = ops.ImageDecoder(device="mixed", output_type=types.RGB)
self.resize = ops.Resize(device="gpu", resize_shorter=256.)
self.cmnp = ops.CropMirrorNormalize(device="gpu",
output_dtype=types.FLOAT,
crop=(224, 224),
image_type=types.RGB,
mean=[0., 0., 0.],
std=[1., 1., 1.])
self.uniform = ops.Uniform(range=(0.0, 1.0))
self.iter = 0
def __init__(self, args):
super(TFRecordDetectionPipeline,
self).__init__(args.batch_size, args.num_workers, 0, 0)
self.input = ops.TFRecordReader(
path=os.path.join(test_dummy_data_path, 'small_coco.tfrecord'),
index_path=os.path.join(test_dummy_data_path,
'small_coco_index.idx'),
features={
'image/encoded':
tfrec.FixedLenFeature((), tfrec.string, ""),
'image/object/class/label':
tfrec.VarLenFeature([1], tfrec.int64, 0),
'image/object/bbox':
tfrec.VarLenFeature([4], tfrec.float32, 0.0),
},
shard_id=0,
num_shards=1,
random_shuffle=False)
self.decode_gpu = ops.ImageDecoder(device="mixed",
output_type=types.RGB)
self.cast = ops.Cast(dtype=types.INT32)
self.box_encoder = ops.BoxEncoder(device="cpu",
criteria=0.5,
anchors=coco_anchors())
def __init__(self, tfrecord_files, tfrecord_index_files, num_shards, shard_id,
dataset_size=None, stick_to_shard=False, pad_last_batch=False,
initial_fill=1024, lazy_init=False, read_ahead=False, prefetch_queue_depth=1,
seed=-1, random_shuffle=False, skip_cached_images=False):
self.tfrecord_files = tfrecord_files
self.tfrecord_index_files = tfrecord_index_files
self.num_shards = num_shards
self.shard_id = shard_id
features = {'image/encoded': tfrec.FixedLenFeature([], tfrec.string, ""),
'image/id': tfrec.FixedLenFeature([], tfrec.int64, -1),
'image/shape': tfrec.FixedLenFeature([3], tfrec.int64, -1),
'image/object/label': tfrec.VarLenFeature([1], tfrec.int64, -1),
'image/object/bbox': tfrec.VarLenFeature([4], tfrec.float32, 0.0),}
# Note that (shuffle_after_epoch, skip_empty, ltrb, ratio, size_threshold) arguments
# are not supported in the TFRecord reader, but some of them can be passed to the TFRecord
# creation script
self.dataset_reader = dali.ops.TFRecordReader(path=self.tfrecord_files,
index_path=self.tfrecord_index_files,
num_shards=self.num_shards,
shard_id=self.shard_id,
stick_to_shard=stick_to_shard,
pad_last_batch=pad_last_batch,
initial_fill=initial_fill,
lazy_init=lazy_init,
read_ahead=read_ahead,
prefetch_queue_depth=prefetch_queue_depth,
seed=seed,
random_shuffle=random_shuffle,
skip_cached_images=skip_cached_images,
features=features)
self._size = dataset_size
self.cast_int32 = dali.ops.Cast(device="cpu", dtype=dali.types.INT32)
def __init__(self, batch_size, num_threads, device_id):
super(TFRecordPipeline, self).__init__(batch_size,
num_threads,
device_id)
self.input = ops.TFRecordReader(path = tfrecord,
index_path = tfrecord_idx,
features = {"image/encoded" : tfrec.FixedLenFeature((), tfrec.string, ""),
'image/filename': tfrec.FixedLenFeature([ ], tfrec.string, ''),
'image/height': tfrec.FixedLenFeature([1], tfrec.int64, -1),
'image/width': tfrec.FixedLenFeature([1], tfrec.int64, -1),
'image/colorspace': tfrec.FixedLenFeature([ ], tfrec.string, ''),
'image/channels': tfrec.FixedLenFeature([1], tfrec.int64, -1),
'image/format': tfrec.FixedLenFeature([ ], tfrec.string, ''),
'image/class/label': tfrec.FixedLenFeature([1], tfrec.int64, -1),
'image/class/synset': tfrec.FixedLenFeature([ ], tfrec.string, ''),
'image/class/text': tfrec.FixedLenFeature([ ], tfrec.string, ''),
'image/object/bbox/xmin': tfrec.VarLenFeature(tfrec.float32, 0.0),
'image/object/bbox/ymin': tfrec.VarLenFeature(tfrec.float32, 0.0),
'image/object/bbox/xmax': tfrec.VarLenFeature(tfrec.float32, 0.0),
'image/object/bbox/ymax': tfrec.VarLenFeature(tfrec.float32, 0.0),
'image/object/bbox/label': tfrec.FixedLenFeature([1], tfrec.int64, -1)})
self.decode = ops.ImageDecoder(device = "mixed", output_type = types.RGB)
self.resize = ops.Resize(device = "gpu", resize_x = 1920.,resize_y=1920.)
self.vert_flip = ops.Flip(device = "gpu", horizontal=0)
self.vert_coin = ops.CoinFlip(probability=0.5)
self.cmnp = ops.CropMirrorNormalize(device = "gpu",
output_dtype = types.FLOAT,
crop = (1920, 1920),
image_type = types.RGB,
mean = [0., 0., 0.],
std = [1., 1., 1.])
self.uniform = ops.Uniform(range = (0.0, 1.0))
self.iter = 0
def __init__(self, batch_size, num_threads, device_id, tfrecords,
idx_paths):
super(ResnetPipeline, self).__init__(batch_size, num_threads,
device_id)
# Transformation operations below.
# From https://docs.nvidia.com/deeplearning/sdk/dali-developer-guide/docs/supported_ops.html
self.input = ops.TFRecordReader(
path=tfrecords,
index_path=idx_paths,
features={
"image/encoded": tfrec.FixedLenFeature([], tfrec.string, ""),
"image/class/label": tfrec.FixedLenFeature([1], tfrec.float32,
0.0),
"image/class/text": tfrec.FixedLenFeature([], tfrec.string,
""),
"image/object/bbox/xmin":
tfrec.VarLenFeature(tfrec.float32, 0.0),
"image/object/bbox/ymin":
tfrec.VarLenFeature(tfrec.float32, 0.0),
"image/object/bbox/xmax":
tfrec.VarLenFeature(tfrec.float32, 0.0),
"image/object/bbox/ymax":
tfrec.VarLenFeature(tfrec.float32, 0.0)
})
self.decode = ops.nvJPEGDecoder(device="mixed",
cache_debug=True,
output_type=types.RGB)
self.resize = ops.Resize(device="gpu",
image_type=types.RGB,
interp_type=types.INTERP_LINEAR,
resize_shorter=256.)
self.cmn = ops.CropMirrorNormalize(device="gpu",
output_dtype=types.FLOAT,
crop=(224, 224),
image_type=types.RGB,
mean=[0., 0., 0.],
std=[1., 1., 1])
self.uniform = ops.Uniform(range=(0.0, 1.0))
self.transpose = ops.Transpose(device="gpu", perm=[0, 3, 1, 2])
self.cast = ops.Cast(device="gpu", dtype=types.INT32)
self.iter = 0
def tfrecord_pipeline(dspath, batch_size, num_threads, device="cpu", device_id=None,
shard_id=0, num_shards=1, reader_name="Reader",
seq=True, chroms=False, chroms_vlog=False, target=True, target_vlog=True, label=False, random_shuffle=True):
pipe = Pipeline(batch_size=batch_size, num_threads=num_threads, device_id=device_id)
feature_description = {}
feature_description["seq"] = tfrec.VarLenFeature(tfrec.float32, -1.0)
feature_description["label"] = tfrec.FixedLenFeature([], tfrec.int64, -1)
feature_description["target"] = tfrec.FixedLenFeature([], tfrec.float32, -1.0)
for ct in dspath["chromatin_tracks"]:
feature_description[ct] = tfrec.VarLenFeature(tfrec.float32, -1.0)
with pipe:
inputs = fn.readers.tfrecord(
name=reader_name,
path=dspath['TFRecord'],
index_path=dspath['TFRecord_idx'],
features=feature_description,
shard_id = shard_id,
num_shards = num_shards,
random_shuffle=random_shuffle,
read_ahead=True,
prefetch_queue_depth=20,
pad_last_batch=True)
if device=="gpu":
inputs['seq'] = inputs['seq'].gpu()
for ct in dspath["chromatin_tracks"]: inputs[ct] = inputs[ct].gpu()
inputs['target'] = inputs['target'].gpu()
inputs['label'] = inputs['label'].gpu()
seqdata = fn.expand_dims(inputs['seq'], axes=1, device=device)
seqdata = fn.reshape(seqdata, shape=(4, -1), device=device)
chromsdata = fn.cat(*[fn.expand_dims(inputs[ct], axes=0, device=device) for ct in dspath["chromatin_tracks"]], axis=0, device=device)
sample = []
if seq: sample.append(seqdata)
if chroms:
if chroms_vlog:
sample.append(log(chromsdata + 1))
else:
sample.append(chromsdata)
if target:
if target_vlog:
sample.append(log(inputs['target'] + 1))
else:
sample.append(inputs['target'])
if label: sample.append(inputs['label'])
pipe.set_outputs(*sample)
return pipe
def __init__(self, batch_size, num_threads, device_id, **kwargs):
super(TFRecordTrain, self).__init__(batch_size, num_threads, device_id)
self.dim = kwargs["dim"]
self.seed = kwargs["seed"]
self.oversampling = kwargs["oversampling"]
self.input = ops.TFRecordReader(
path=kwargs["tfrecords"],
index_path=kwargs["tfrecords_idx"],
features={
"X_shape": tfrec.FixedLenFeature([self.dim + 1], tfrec.int64, 0),
"Y_shape": tfrec.FixedLenFeature([self.dim + 1], tfrec.int64, 0),
"X": tfrec.VarLenFeature([], tfrec.float32, 0.0),
"Y": tfrec.FixedLenFeature([], tfrec.string, ""),
"fname": tfrec.FixedLenFeature([], tfrec.string, ""),
},
num_shards=kwargs["gpus"],
shard_id=device_id,
random_shuffle=True,
pad_last_batch=True,
read_ahead=True,
seed=self.seed,
)
self.patch_size = kwargs["patch_size"]
self.crop_shape = types.Constant(np.array(self.patch_size), dtype=types.INT64)
self.crop_shape_float = types.Constant(np.array(self.patch_size), dtype=types.FLOAT)
self.layout = "CDHW" if self.dim == 3 else "CHW"
self.axis_name = "DHW" if self.dim == 3 else "HW"
def __init__(self, batch_size, num_threads, device_id, **kwargs):
super(TFRecordTest, self).__init__(batch_size, num_threads, device_id)
self.input = ops.TFRecordReader(
path=kwargs["tfrecords"],
index_path=kwargs["tfrecords_idx"],
features={
"X_shape": tfrec.FixedLenFeature([4], tfrec.int64, 0),
"X": tfrec.VarLenFeature([], tfrec.float32, 0.0),
"fname": tfrec.FixedLenFeature([], tfrec.string, ""),
},
shard_id=device_id,
num_shards=kwargs["gpus"],
read_ahead=True,
random_shuffle=False,
pad_last_batch=True,
)
def tfrecord_pipe_empty_fields(path, index_path):
inputs = fn.readers.tfrecord(path=path,
index_path=index_path,
features={
"image/encoded":
tfrec.FixedLenFeature((), tfrec.string,
""),
"image/class/label":
tfrec.FixedLenFeature([1], tfrec.int64,
-1),
"does/not/exists":
tfrec.VarLenFeature(tfrec.int64, -1),
"does/not/exists/as/well":
tfrec.FixedLenFeature([1], tfrec.float32,
.0)
})
return inputs["image/encoded"], inputs["does/not/exists"], inputs[
"does/not/exists/as/well"]
def __init__(self, batch_size, num_threads, device_id, **kwargs):
super(TFRecordBenchmark, self).__init__(batch_size, num_threads, device_id)
self.dim = kwargs["dim"]
self.input = ops.TFRecordReader(
path=kwargs["tfrecords"],
index_path=kwargs["tfrecords_idx"],
features={
"X_shape": tfrec.FixedLenFeature([self.dim + 1], tfrec.int64, 0),
"Y_shape": tfrec.FixedLenFeature([self.dim + 1], tfrec.int64, 0),
"X": tfrec.VarLenFeature([], tfrec.float32, 0.0),
"Y": tfrec.FixedLenFeature([], tfrec.string, ""),
"fname": tfrec.FixedLenFeature([], tfrec.string, ""),
},
shard_id=device_id,
num_shards=kwargs["gpus"],
read_ahead=True,
)
self.patch_size = kwargs["patch_size"]
self.layout = "CDHW" if self.dim == 3 else "CHW"
def get_dali_pipeline(tfrec_filenames,
tfrec_idx_filenames,
height,
width,
shard_id,
num_gpus,
dali_cpu=True,
training=True):
inputs = fn.readers.tfrecord(path=tfrec_filenames,
index_path=tfrec_idx_filenames,
random_shuffle=training,
shard_id=shard_id,
num_shards=num_gpus,
initial_fill=10000,
features={
'image/encoded':
tfrec.FixedLenFeature((), tfrec.string,
""),
'image/class/label':
tfrec.FixedLenFeature([1], tfrec.int64,
-1),
'image/class/text':
tfrec.FixedLenFeature([], tfrec.string,
''),
'image/object/bbox/xmin':
tfrec.VarLenFeature(tfrec.float32, 0.0),
'image/object/bbox/ymin':
tfrec.VarLenFeature(tfrec.float32, 0.0),
'image/object/bbox/xmax':
tfrec.VarLenFeature(tfrec.float32, 0.0),
'image/object/bbox/ymax':
tfrec.VarLenFeature(tfrec.float32, 0.0)
})
decode_device = "cpu" if dali_cpu else "mixed"
resize_device = "cpu" if dali_cpu else "gpu"
if training:
images = fn.decoders.image_random_crop(
inputs["image/encoded"],
device=decode_device,
output_type=types.RGB,
random_aspect_ratio=[0.75, 1.25],
random_area=[0.05, 1.0],
num_attempts=100)
images = fn.resize(images,
device=resize_device,
resize_x=width,
resize_y=height)
else:
images = fn.decoders.image(inputs["image/encoded"],
device=decode_device,
output_type=types.RGB)
# Make sure that every image > 224 for CropMirrorNormalize
images = fn.resize(images, device=resize_device, resize_shorter=256)
images = fn.crop_mirror_normalize(images.gpu(),
dtype=types.FLOAT,
crop=(height, width),
mean=[123.68, 116.78, 103.94],
std=[58.4, 57.12, 57.3],
output_layout="HWC",
mirror=fn.random.coin_flip())
labels = inputs["image/class/label"].gpu()
labels -= 1 # Change to 0-based (don't use background class)
return images, labels
def __init__(self,
tfrecords,
batch_size,
target_size,
preproc_param,
num_threads,
num_shards,
device_ids,
training=False):
Pipeline.__init__(self,
batch_size=batch_size,
num_threads=num_threads,
device_id=device_ids,
prefetch_queue_depth=num_threads,
seed=42,
exec_async=False,
exec_pipelined=False)
DaliPipeline.__init__(self,
target_size=target_size,
preproc_param=preproc_param,
training=training)
tfrecords_idx = [tfrecord + "_idx" for tfrecord in tfrecords]
for tfrecord, tfrecord_idx in zip(tfrecords, tfrecords_idx):
if os.path.exists(tfrecord_idx):
continue
call(["tfrecord2idx", tfrecord, tfrecord + "_idx"])
self.length = sum([len(open(f).readlines()) for f in tfrecords_idx])
self.input = ops.TFRecordReader(
path=tfrecords,
index_path=tfrecords_idx,
features={
'image/height':
tfrec.FixedLenFeature([1], tfrec.int64, -1),
'image/width':
tfrec.FixedLenFeature([1], tfrec.int64, -1),
'image/encoded':
tfrec.FixedLenFeature((), tfrec.string, ""),
'image/format':
tfrec.FixedLenFeature((), tfrec.string, ""),
'image/object/bbox/xmin':
tfrec.VarLenFeature(tfrec.float32, 0.0),
'image/object/bbox/ymin':
tfrec.VarLenFeature(tfrec.float32, 0.0),
'image/object/bbox/xmax':
tfrec.VarLenFeature(tfrec.float32, 0.0),
'image/object/bbox/ymax':
tfrec.VarLenFeature(tfrec.float32, 0.0),
'image/object/class/text':
tfrec.FixedLenFeature([], tfrec.string, ''),
'image/object/class/label':
tfrec.VarLenFeature(tfrec.int64, -1)
},
num_shards=num_shards,
random_shuffle=training)
self.training = training
self.cat = dalitorch.TorchPythonFunction(
function=lambda l, t, r, b: torch.cat([l, t, r, b]).view(
4, -1).permute(1, 0)) #[l*w,t*h,r*w,b*h], [l,t,r,b]
self.cast = ops.Cast(dtype=types.DALIDataType.INT32)
