Implementation:NVIDIA DALI DALIDatasetOp
| Knowledge Sources | |
|---|---|
| Domains | TensorFlow_Integration, Data_Pipeline |
| Last Updated | 2026-02-08 16:00 GMT |
Overview
Implements the DALIDataset TensorFlow custom op that integrates DALI data pipelines as tf.data.Dataset sources for both CPU and GPU devices.
Description
This file provides the full implementation of the DALIDatasetOp TensorFlow custom operation, which wraps a serialized DALI pipeline into a tf.data.Dataset compatible interface. The implementation consists of two major nested classes: Dataset (derived from DatasetBase) and Iterator (derived from DatasetIterator). The Dataset class manages the pipeline definition, output shapes and dtypes, device placement, and serialization to TensorFlow's graph definition format. It also supports input datasets that can be fed into DALI External Source nodes.
The Iterator class handles the runtime lifecycle of the DALI pipeline, including initialization, prefetching, input feeding, output production, checkpointing (save/restore), and teardown. It implements a state machine for tracking input dataset exhaustion with three states: in_progress, stop_pending, and stop_signaled. The iterator feeds input batches to DALI's External Source operators in either per-sample or batched mode, manages a queue of alive batches for memory lifetime, and copies DALI pipeline outputs into TensorFlow tensors with appropriate stream synchronization for GPU operations.
The file also contains the TensorFlow op registration (REGISTER_OP("DALIDataset")) and kernel builder registrations for both CPU and GPU devices. The op registration defines all attributes including pipeline configuration, input/output specifications, and device mismatch handling. Version-specific preprocessor guards ensure compatibility across TensorFlow versions from 1.15 through 2.x, handling API differences in split providers, checkpointing, cardinality, and status codes.
Usage
This implementation is compiled as part of the DALI TensorFlow plugin shared library and is used whenever DALIDataset is instantiated in TensorFlow Python code via dali_tf_plugin. It provides the core bridge between DALI's GPU-accelerated data processing pipelines and TensorFlow's tf.data API.
Code Reference
Source Location
- Repository: NVIDIA_DALI
- File: dali_tf_plugin/dali_dataset_op.cc
- Lines: 1-1224
Signature
class DALIDatasetOp::Dataset : public DatasetBase {
public:
explicit Dataset(OpKernelContext *context, const PipelineDef pipeline_def,
const Inputs &inputs, const InputAttrs &input_attrs,
const std::vector<PartialTensorShape> &shapes, const DataTypeVector &dtypes,
const bool is_gpu_device, const bool fail_on_device_mismatch);
std::unique_ptr<IteratorBase> MakeIteratorInternal(const string &prefix) const override;
const DataTypeVector &output_dtypes() const override;
const std::vector<PartialTensorShape> &output_shapes() const override;
};
class DALIDatasetOp::Dataset::Iterator : public DatasetIterator<Dataset> {
public:
explicit Iterator(const Params ¶ms,
dali::c_api::PipelineHandle pipeline_handle,
bool enable_memory_stats = false);
Status Initialize(IteratorContext *context) override;
Status GetNextInternal(IteratorContext *context, std::vector<Tensor> *out_tensors,
bool *end_of_sequence) override;
};
void DALIDatasetOp::MakeDataset(OpKernelContext *context, DatasetBase **output);
REGISTER_OP("DALIDataset")
.Input("input_datasets: N * variant")
.Output("handle: variant")
.Attr("pipeline: string")
.Attr("batch_size: int")
.Attr("num_threads: int")
.Attr("device_id: int")
.Attr("output_shapes: list(shape) >= 1")
.Attr("output_dtypes: list({bool, half, float, uint8, uint16, uint32, uint64, int8, int16, int32, int64}) >= 1");
Import
#include "dali_tf_plugin/dali_dataset.h"
#include "dali_tf_plugin/dali_helper.h"
#include "dali/dali.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| input_datasets | N * variant | No | Zero or more TF datasets to feed into DALI External Source nodes |
| pipeline | string | Yes | Serialized DALI pipeline definition |
| batch_size | int | Yes | Maximum batch size for the DALI pipeline |
| num_threads | int | Yes | Number of CPU threads for the DALI pipeline |
| device_id | int | Yes | GPU device ID (-1 for CPU-only) |
| exec_separated | bool | No | Whether to use separated executor (default false) |
| exec_dynamic | bool | No | Whether to use dynamic executor (default false) |
| prefetch_queue_depth | int | Yes | Depth of the uniform prefetch queue |
| cpu_prefetch_queue_depth | int | Yes | CPU prefetch queue depth (separated executor) |
| gpu_prefetch_queue_depth | int | Yes | GPU prefetch queue depth (separated executor) |
| output_shapes | list(shape) | Yes | Expected output tensor shapes |
| output_dtypes | list(type) | Yes | Expected output tensor data types |
| fail_on_device_mismatch | bool | No | Raise error on TF/DALI device mismatch (default true) |
| input_names | list(string) | No | Names of DALI External Source inputs |
| input_layouts | list(string) | No | Data layouts for each input |
| input_batched | list(int) | No | Whether each input is batched (1) or per-sample (0) |
Outputs
| Name | Type | Description |
|---|---|---|
| handle | variant | A tf.data.Dataset handle that yields tensors matching output_shapes and output_dtypes |
Usage Examples
Creating a DALIDataset in Python
import nvidia.dali as dali
from nvidia.dali.plugin.tf import DALIDataset
import tensorflow as tf
@dali.pipeline_def(batch_size=32, num_threads=4, device_id=0)
def my_pipeline():
images, labels = dali.fn.readers.file(file_root="/data/images", name="Reader")
images = dali.fn.decoders.image(images, device="mixed")
images = dali.fn.resize(images, size=(224, 224))
return images, labels
pipe = my_pipeline()
pipe.build()
dataset = DALIDataset(
pipeline=pipe,
output_shapes=((32, 224, 224, 3), (32, 1)),
output_dtypes=(tf.uint8, tf.int32),
device_id=0,
batch_size=32,
num_threads=4,
)
for batch in dataset.take(10):
images, labels = batch
# process images and labels