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)

创建与 tf.data.Dataset 兼容的 DALIDataset，它来自 DALI pipeline。它支持 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) – 定义要执行的数据处理的 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 的批处理大小。

• 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 更能抵抗每个批次的不均匀执行时间，但它也会消耗更多内存用于内部缓冲区。该值将与 exec_separated 设置为 False 一起使用。

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

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

返回类型:

基于 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 插件

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

参数:

• serialized_pipeline – 一个 string。

• 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 变体。此数据集增加了对输入 tf.data.Datasets 的支持。对输入 tf.data.Datasets 的支持仅适用于 TensorFlow 2.4.1 及更高版本。

输入数据集规范

每个输入数据集都必须映射到一个 external_source() 算子，该算子将表示 DALI pipeline 的输入。在 pipeline 中，输入由 external_source() 的 name 参数表示。输入数据集必须通过 input_datasets 字典参数从该 name 映射到数据集对象来提供给 DALIDatasetWithInputs。

每样本和批处理模式

输入数据集可以在每样本模式或批处理模式下运行。

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

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

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

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

带有 source 参数的外部源

此实验性 DALIDataset 接受带有 external_source() 节点的 pipeline，这些节点指定了 source 参数。在这种情况下，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)
}

可以作为 input_datasets 传递给如下 Pipeline

@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, optional, default = None) –

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

  布局必须相同。

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

• batch (bool, optional, default = False) –

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

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

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