TensorFlow 插件 API 参考

class nvidia.dali.plugin.tf.DALIDataset(pipeline, output_dtypes=None, output_shapes=None, fail_on_device_mismatch=True, *, input_datasets=None, batch_size=1, num_threads=4, device_id=0, exec_separated=False, exec_dynamic=False, prefetch_queue_depth=2, cpu_prefetch_queue_depth=2, gpu_prefetch_queue_depth=2, dtypes=None, shapes=None)

创建与来自 DALI pipeline 的 tf.data.Dataset 兼容的 DALIDataset。它支持 TensorFlow 1.15 和 2.x 系列。

DALIDataset 可以放置在 CPU 和 GPU 上。

请记住，TensorFlow 默认分配几乎所有可用的设备内存。这可能会由于内存不足而在 DALI 中引起错误。关于如何更改此行为，请查看 TensorFlow 文档，因为它可能因您的用例而异。

警告

大多数 TensorFlow Dataset 只有 CPU 变体。要通过其他 TensorFlow dataset 处理 GPU 放置的 DALIDataset，您需要首先使用显式的 tf.data.experimental.copy_to_device 将其复制回 CPU - 从 CPU 到 GPU 再到 CPU 的往返可能会大大降低性能，因此不建议这样做。

此外，建议不要在 DALIDataset 之后使用像 repeat() 或类似的 dataset，这可能会干扰 DALI 内存分配和预取。

参数:

• pipeline (nvidia.dali.Pipeline) – 定义要执行的数据处理。

• output_dtypes (tf.DType 或 tf.DType 的元组，默认值 = None) – 预期输出类型

• output_shapes (形状的元组，可选，默认值 = None) – 预期输出形状。如果提供，则必须与 output_dtypes 的arity 匹配。当设置为 None 时，DALI 将自行推断形状。单个形状也可以设置为 None 或包含 None 以指示未知维度。如果指定，则必须与从 DALI Pipeline 返回的形状以及 batch_size 参数兼容，这将是返回张量的最外层维度。在 batch_size = 1 的情况下，可以在形状中省略它。DALI Dataset 将尝试通过从从 Pipeline 获取的形状中挤压 1 大小的维度来匹配请求的形状。

• fail_on_device_mismatch (bool，可选，默认值 = True) –

  当设置为 True 时，将执行运行时检查以确保 DALI 设备和 TF 设备都是 CPU 或都是 GPU。在某些情况下，此检查可能不准确。当设置为

  False 将跳过检查，但会打印额外的日志以检查设备。请记住

  这可能会允许在工作流程中隐藏 GPU 到 CPU 的复制，并影响性能。

• batch_size (int，可选，默认值 = 1) – pipeline 的 batch size。

• num_threads (int，可选，默认值 = 4) – pipeline 使用的 CPU 线程数。

• device_id (int，可选，默认值 = 0) – pipeline 使用的 GPU 的 ID。此参数的 None 值表示 DALI 不应使用 GPU 或 CUDA 运行时。这会将 pipeline 限制为仅 CPU 运算符，但允许它在任何具有 CPU 功能的机器上运行。

• exec_separated (bool，可选，默认值 = False) – 是否以允许 CPU 和 GPU 计算重叠的方式执行 pipeline，通常会导致更快的执行速度，但更大的内存消耗。此标志与 exec_dymamic 不兼容。

• exec_dynamic (bool，可选，默认值 = False) – 是否使用动态执行器执行 pipeline，这允许灵活地混合 CPU 和 GPU 运算符，并实现积极的内存重用。此标志与 exec_separated 不兼容。

• prefetch_queue_depth (int，可选，默认值 = 2) – 执行器队列的深度。更深的队列使 DALI 更能抵抗每个 batch 不均匀的执行时间，但它也会消耗更多内存用于内部缓冲区。该值将与设置为 False 的 exec_separated 一起使用。

• cpu_prefetch_queue_depth (int，可选，默认值 = 2) – 执行器 cpu 队列的深度。更深的队列使 DALI 更能抵抗每个 batch 不均匀的执行时间，但它也会消耗更多内存用于内部缓冲区。该值将与设置为 True 的 exec_separated 一起使用。

• gpu_prefetch_queue_depth (int，可选，默认值 = 2) – 执行器 gpu 队列的深度。更深的队列使 DALI 更能抵抗每个 batch 不均匀的执行时间，但它也会消耗更多内存用于内部缓冲区。该值将与设置为 True 的 exec_separated 一起使用。

返回类型:

基于 DALI pipeline 且与 tf.data.Dataset API 兼容的 DALIDataset 对象。

nvidia.dali.plugin.tf.DALIIterator()

TF 插件包装器

此运算符的工作方式与 DALI TensorFlow 插件相同，但例外的是它也接受 Pipeline 对象作为输入，这些对象在内部序列化。有关更多信息，请参见 nvidia.dali.plugin.tf.DALIRawIterator()。

nvidia.dali.plugin.tf.DALIIteratorWrapper(pipeline=None, serialized_pipeline=None, sparse=[], shapes=[], dtypes=[], batch_size=-1, prefetch_queue_depth=2, exec_dynamic=False, **kwargs)

TF 插件包装器

此运算符的工作方式与 DALI TensorFlow 插件相同，但例外的是它也接受 Pipeline 对象作为输入，这些对象在内部序列化。有关更多信息，请参见 nvidia.dali.plugin.tf.DALIRawIterator()。

nvidia.dali.plugin.tf.DALIRawIterator()

DALI TensorFlow 插件

从序列化的 pipeline 创建 DALI pipeline，从 serialized_pipeline 参数获取。shapes 必须与对应的 DALI Pipeline 输出张量形状的形状匹配。dtypes 必须与对应的 DALI Pipeline 输出张量类型的类型匹配。

参数:

• serialized_pipeline – 一个 字符串。

• shapes – 形状列表（每个形状都是 tf.TensorShape 或 ints 列表），长度 >= 1。

• dtypes – tf.DTypes 列表，来自：tf.half, tf.float32, tf.uint8, tf.int16, tf.int32, tf.int64，长度 >= 1。

• num_threads – 一个可选的 int。默认为 -1。

• device_id – 一个可选的 int。默认为 -1。

• exec_separated – 一个可选的 bool。默认为 False。

• exec_dynamic – 一个可选的 bool。默认为 False。

• gpu_prefetch_queue_depth – 一个可选的 int。默认为 2。

• cpu_prefetch_queue_depth – 一个可选的 int。默认为 2。

• sparse – 一个可选的 bools 列表。默认为 []。

• batch_size – 一个可选的 int。默认为 -1。

• enable_memory_stats – 一个可选的 bool。默认为 False。

• name – 操作的名称（可选）。

返回值:

类型为 dtypes 的 Tensor 对象列表。

请记住，TensorFlow 默认分配几乎所有可用的设备内存。这可能会由于内存不足而在 DALI 中引起错误。关于如何更改此行为，请查看 TensorFlow 文档，因为它可能因您的用例而异。

nvidia.dali.plugin.tf.dataset_compatible_tensorflow()

如果当前 TensorFlow 版本与 DALIDataset 兼容，则返回 True。

nvidia.dali.plugin.tf.dataset_distributed_compatible_tensorflow()

如果当前 TensorFlow 版本的 tf.distribute API 与 DALIDataset 兼容，则返回 True。

nvidia.dali.plugin.tf.dataset_inputs_compatible_tensorflow()

如果当前 TensorFlow 版本与 experimental.DALIDatasetWithInputs 兼容并且输入 Dataset 可以与 DALI 一起使用，则返回 True。

nvidia.dali.plugin.tf.dataset_options()

nvidia.dali.plugin.tf.serialize_pipeline(pipeline)

实验性

nvidia.dali.plugin.tf.experimental.DALIDatasetWithInputs(pipeline, output_dtypes=None, output_shapes=None, fail_on_device_mismatch=True, *, input_datasets=None, batch_size=1, num_threads=4, device_id=0, exec_separated=False, exec_dynamic=False, prefetch_queue_depth=2, cpu_prefetch_queue_depth=2, gpu_prefetch_queue_depth=2, dtypes=None, shapes=None)

实验性变体 DALIDataset。此 dataset 添加了对输入 tf.data.Datasets 的支持。对输入 tf.data.Datasets 的支持仅适用于 TensorFlow 2.4.1 及更高版本。

输入 dataset 规范

每个输入 dataset 都必须映射到一个 external_source() 运算符，该运算符将表示 DALI pipeline 的输入。在 pipeline 中，输入表示为 external_source() 的 name 参数。输入 dataset 必须作为从该 name 到 dataset 对象的映射通过 DALIDatasetWithInputs 的 input_datasets 字典参数提供。

每样本和 batch 模式

输入 dataset 可以以每样本模式或 batch 模式运行。

在逐样本模式下，源数据集产生的值被解释为单个样本。批次维度不存在。例如，一个 640x480 RGB 图像的形状为 [480, 640, 3]。

在批次模式下，源数据集生成的张量被解释为批次，并增加了一个额外的外部维度来表示批次中的样本。例如，包含十个 640x480 RGB 图像的批次的形状为 [10, 480, 640, 3]。

在这两种情况下（逐样本和批次模式），这些输入的布局都应表示为 “HWC”。

在逐样本模式下，DALIDataset 将查询输入数据集 batch_size 次，以构建将馈送到 DALI Pipeline 的批次。在逐样本模式下，输入数据集生成的每个样本可以具有不同的形状，但维度数量和布局必须保持不变。

带有 source 参数的外部源

这个实验性的 DALIDataset 接受带有 external_source() 节点且指定了 source 参数的 Pipeline。在这种情况下，source 将自动转换为适当的 tf.data.Dataset.from_generator 数据集，并具有正确的放置和 tf.data.experimental.copy_to_device 指令。

这些节点也可以在逐样本或批次模式下工作。批次模式下的数据必须是密集的、均匀的张量（每个样本具有相同的维度）。仅接受 CPU 数据。

这使得 TensorFlow DALIDataset 可以与大多数已经指定了外部源 source 的 Pipeline 一起工作。

警告

这个类是实验性的，其 API 可能会在没有通知的情况下更改。

注意

不支持将外部源节点的 num_outputs 指定为任何数字 - 这意味着不支持具有多个（元组）输出的回调。

注意

不支持外部源 cycle 策略 'raise' - 数据集不可重启。

注意

外部源 cuda_stream 参数被忽略 - source 应该返回 CPU 数据，并且 tf.data.Dataset 输入在内部处理。

注意

外部源 use_copy_kernel 和 blocking 参数被忽略。

注意

当与 DALI Dataset 一起使用时，在定义 pipeline 时，在外部源节点上设置 no_copy 被认为是无操作。 no_copy 选项在内部处理，并在可能的情况下自动启用。

注意

不支持通过 source 提供的外部源回调的并行执行。回调通过 TensorFlow tf.data.Dataset.from_generator 执行 - parallel 和 prefetch_queue_depth 参数被忽略。

该运算符添加了额外的参数，以补充 DALIDataset 支持的参数。

参数:

input_datasets (dict[str, tf.data.Dataset] 或)

dict[str, nvidia.dali.plugin.tf.experimental.Input] DALI Pipeline 的输入数据集。它必须作为字典提供，从 External Source 节点的名称映射到数据集对象或 Input() 包装器。

例如

{
    'tensor_input': tf.data.Dataset.from_tensors(tensor).repeat(),
    'generator_input': tf.data.Dataset.from_generator(some_generator)
}

可以作为 Pipeline 的 input_datasets 传递，例如

@pipeline_def
def external_source_pipe():
    input_0 = fn.external_source(name='tensor_input')
    input_1 = fn.external_source(name='generator_input')
    return fn.resize(input_1, resize_x=input_0)

直接使用 tf.data.Dataset 的条目，例如

{
    'input': tf.data.Dataset.from_tensors(tensor)
}

等效于使用 nvidia.dali.plugin.tf.experimental.Input 的以下规范

{
    'input' : nvidia.dali.plugin.tf.experimental.Input(
                  dataset=tf.data.Dataset.from_tensors(tensor),
                  layout=None,
                  batch=False)
}

这意味着直接指定为 tf.data.Dataset 的输入被视为样本输入。

警告

输入数据集必须与 DALIDatasetWithInputs 放置在同一设备上。如果输入具有不同的放置位置（例如，输入放置在 CPU 上，而 DALIDatasetWithInputs 放置在 GPU 上），则必须首先将带有 GPU 参数的 tf.data.experimental.copy_to_device 应用于输入。

nvidia.dali.plugin.tf.experimental.Input(dataset, *, layout=None, batch=False)

传递给 DALIDataset 的输入的包装器。允许传递额外的选项，这些选项可以覆盖 Python Pipeline 对象中 External Source 节点中指定的一些选项。传递 None 表示应在 pipeline 定义中查找该值。

参数:

• dataset (tf.data.Dataset) – 用作输入的 dataset

• layout (str, 可选, 默认 = None) –

  输入的布局。如果为 None，则布局将从 Python Pipeline 对象中相应的 External Source 节点获取。如果两者都提供，

  则布局必须相同。

  如果两者都未提供，将使用空布局。

• batch (bool, 可选, 默认 = False) –

  给定输入的批次模式。如果为 None，则批次模式将从 Python Pipeline 对象中相应的 External Source 节点获取。

  如果 batch = False，则输入数据集被视为样本输入。

  如果 batch = True，则输入数据集应返回批次。