Implementation:NVIDIA DALI DALIDatasetOp Header
| Knowledge Sources | |
|---|---|
| Domains | TensorFlow_Integration, Data_Pipeline |
| Last Updated | 2026-02-08 16:00 GMT |
Overview
Declares the DALIDatasetOp class, its nested data structures for pipeline definitions, inputs, and input attributes, along with the forward declaration of the inner Dataset class.
Description
This header file defines the DALIDatasetOp class within the dali_tf_impl namespace. The class extends TensorFlow's DatasetOpKernel to implement a custom dataset op that wraps a DALI pipeline. The constructor extracts all configuration attributes from the OpKernelConstruction context, including the pipeline definition parameters, input attributes, output shapes and dtypes, and device mismatch behavior.
The header defines three key internal data structures: PipelineDef holds the serialized pipeline string and execution parameters (batch size, threads, device ID, executor modes, and prefetch queue depths); Inputs holds a vector of pointers to upstream DatasetBase objects that feed into DALI External Source nodes; and InputAttrs holds the static attributes describing each input (names, layouts, and batched flags). The InputDescs struct combines both Inputs and InputAttrs through multiple inheritance for convenient passing as a single parameter.
The class declares string constants for all TensorFlow attribute keys (e.g., kPipeline, kBatchSize, kInputNames, kOutputShapes), helper methods for extracting attributes from the construction context (FillPipelineDef, FillInputs, FillInputAttrs, ValidateInputs), and the MakeDataset virtual method override. The forward-declared inner Dataset class is implemented in the corresponding dali_dataset_op.cc file. The entire declaration is guarded by TensorFlow version checks requiring TF >= 1.15 or TF 2.x.
Usage
Include this header in the dali_dataset_op.cc implementation file to access the DALIDatasetOp class declaration and its nested types. This header provides the interface contract between the op kernel registration and the dataset/iterator implementation.
Code Reference
Source Location
- Repository: NVIDIA_DALI
- File: dali_tf_plugin/dali_dataset.h
- Lines: 1-144
Signature
namespace dali_tf_impl {
class DALIDatasetOp : public tensorflow::data::DatasetOpKernel {
public:
explicit DALIDatasetOp(tensorflow::OpKernelConstruction* context);
void MakeDataset(tensorflow::OpKernelContext* context,
tensorflow::data::DatasetBase** output) override;
private:
struct PipelineDef {
std::string pipeline;
int batch_size;
int num_threads;
int device_id;
bool exec_separated;
bool exec_dynamic;
int prefetch_queue_depth;
int cpu_prefetch_queue_depth;
int gpu_prefetch_queue_depth;
bool enable_memory_stats;
};
struct Inputs {
std::vector<tensorflow::data::DatasetBase *> inputs;
};
struct InputAttrs {
std::vector<std::string> input_names;
std::vector<std::string> input_layouts;
std::vector<int> input_batched;
};
struct InputDescs : Inputs, InputAttrs {
InputDescs(const Inputs& inputs, const InputAttrs& input_attrs);
};
void FillPipelineDef(tensorflow::OpKernelConstruction* context, PipelineDef& def);
void FillInputs(tensorflow::OpKernelContext *context, Inputs &def);
void FillInputAttrs(tensorflow::OpKernelConstruction* context, InputAttrs& def);
void ValidateInputs(tensorflow::OpKernelContext* context, Inputs& inputs,
InputAttrs& input_attrs);
class Dataset;
};
} // namespace dali_tf_impl
Import
#include "dali_tf_plugin/dali_dataset.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| context | tensorflow::OpKernelConstruction* | Yes | TensorFlow op construction context providing attribute access |
Outputs
| Name | Type | Description |
|---|---|---|
| DALIDatasetOp instance | object | Fully constructed dataset op kernel with pipeline configuration, input attributes, output shapes/dtypes, and device settings |
| DatasetBase** | pointer | Through MakeDataset, produces the Dataset object that implements the tf.data.Dataset interface |
Usage Examples
Including the Header
// In dali_dataset_op.cc
#include "dali_tf_plugin/dali_dataset.h"
namespace dali_tf_impl {
// Implement the nested Dataset class
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);
// ... implementation
};
void DALIDatasetOp::MakeDataset(OpKernelContext *context, DatasetBase **output) {
Inputs inputs;
FillInputs(context, inputs);
ValidateInputs(context, inputs, input_attrs_);
*output = new Dataset(context, pipeline_def_, inputs, input_attrs_,
shapes_, dtypes_, is_gpu_device_, fail_on_device_mismatch_);
}
} // namespace dali_tf_impl