Replacing the Data Loader

The DALI pipeline for StreamPETR is created using the pipeline described in the StreamPETR Pipeline example (including the pipeline itself, the custom processing step, and the custom post-processing function). It also uses the NuScenes loading implementation described in Data Loading for NuScenes.

The goal is to replace the original PyTorch DataLoader (used for the training data loading) with the DALI pipeline. Note that we focus on the training data loading, and do not replace the validation data loading.

This page describes the changes needed to the training script to enable the use of the DALI pipeline.

See the Evaluation for the runtime and CPU usage gains from this integration and instructions for comparing it with the PyTorch data loader.

Note

This demo re-uses the StreamPETR pipeline setup already described in the StreamPETR Pipeline Setup example, and assumes that the pipeline_setup module is available in the Python path. That example primarily describes the design and construction of the DALI pipeline. Please make sure that you are familiar with it before continuing here. This page focuses on the remaining integration work needed to use the completed pipeline inside the StreamPETR training implementation.

Main Training Script

Environment and Paths

Before creating the pipeline, the script configures the Python path to include the DALI example module and project roots needed by the demo.

demos/streampetr/tools/train.py

# ===== Setting up the Python Path (including DALI pipeline) =====

import sys
sys.path.append("/streampetr")
#@NOTE: Make the DALI pipeline example code available.
sys.path.append("/ACCV-Lab/packages/dali_pipeline_framework/examples")

# ===== End of Setting up the Python Path (including DALI pipeline) =====

DALI Pipeline Setup

The major change is the creation of the DALI pipeline instead of a PyTorch Dataset. Both options are supported in our modified implementation, and can be enabled by setting the use_dali_pipeline flag to True in the used configuration file.

demos/streampetr/tools/train.py

    # ===== Get Dataset or DALI Output Iterator =====
    #@NOTE: The if-branch is for the DALI pipeline, the else-branch is for the original PyTorch Dataset.
    if cfg.data_loading_config.use_dali_pipeline:

        #@NOTE 
        # Import the StreamPETR pipeline setup from the DALI pipeline example (and notify user if it is not 
        # found).
        try:
            from pipeline_setup import stream_petr_pipeline
        except ImportError:
            raise ImportError(
                "Could not import `stream_petr_pipeline` from `pipeline_setup`. When running with the DALI "
                "pipeline, please make sure that the `pipeline_setup` module from the "
                "dali_pipeline_framework example (packages/dali_pipeline_framework/examples) is in the python path, as "
                "this demo re-uses the StreamPETR pipeline setup showcased there."
            )

        #@NOTE
        # Use the same dataset root for DALI training and the PyTorch validation dataset. The NuScenes
        # version is specified separately.
        nuscenes_base_dir = cfg.data.train.data_root
        nuscenes_version = cfg.data_loading_config.nuscenes_version

        num_batches_in_epoch = cfg.num_iters_per_epoch

        #@NOTE 
        # Create the DALI pipeline. Note that the returned object is a 
        # `dali_pipeline_framework.DALIStructuredOutputIterator` object, which can be used as a drop-in replacement for 
        # a PyTorch DataLoader. This means that if DALI is used, one needs to skip the creation of a PyTorch 
        # DataLoader from this object, and use it directly instead. This is done in 
        # `projects/mmdet3d_plugin/core/apis/mmdet_train.py`.
        dali_iterator = stream_petr_pipeline.setup_dali_pipeline_stream_petr_train(
            nuscenes_base_dir, nuscenes_version, num_batches_in_epoch, config=cfg.dali_preproc_config, 
            batch_size=cfg.data.samples_per_gpu,
            **get_dali_pipeline_distributed_args(
                local_rank, rank, world_size,
                cfg.data_loading_config.get('replicate_dataset_per_gpu', False),
            ),
        )
        
        #@NOTE 
        # Store as the only element in the `datasets` list. In case DALI is not used, the same variable will 
        # contain the actual PyTorch Dataset (also as a single-element list), see else-branch below.
        datasets = [dali_iterator]
        
        #@NOTE 
        # The classes can be obtrained from the Dataset object if DALI is not used. Here, we manually
        # set the classes instead. 
        classes = [
            'car', 'truck', 'construction_vehicle', 'bus', 'trailer', 'barrier', 'motorcycle', 'bicycle', 
            'pedestrian', 'traffic_cone'
        ]

    else:

        #@NOTE: This corresponds to the original implementation, where a PyTorch Dataset is created and used.

        datasets = [build_dataset(cfg.data.train)]
        if len(cfg.workflow) == 2:
            val_dataset = copy.deepcopy(cfg.data.val)
            # in case we use a dataset wrapper
            if 'dataset' in cfg.data.train:
                val_dataset.pipeline = cfg.data.train.dataset.pipeline
            else:
                val_dataset.pipeline = cfg.data.train.pipeline
            # set test_mode=False here in deep copied config
            # which do not affect AP/AR calculation later
            # refer to https://mmdetection3d.readthedocs.io/en/latest/tutorials/customize_runtime.html#customize-workflow  # noqa
            val_dataset.test_mode = False
            datasets.append(build_dataset(val_dataset))
        
        #@NOTE 
        # Extract the classes here, to obtain the `classes` variable, equivalent to the one created if using
        # DALI. This is done to avoid calling `datasets[0].CLASSES` below, which would not work DALI is used.
        classes = datasets[0].CLASSES
    # ===== End of Getting Dataset or DALI Output Iterator =====

Checkpoint Configuration

The setting of the checkpoint meta information is slightly adjusted to use the classes variable instead of dataset[0].CLASSES (as the latter is not available when DALI is used).

demos/streampetr/tools/train.py

    # ===== Checkpoint Configuration =====
    #@NOTE
    # Note that here, we use the previously obtained `classes` variable. In the original implementation, 
    # `dataset[0].CLASSES` is used instead.
    if cfg.checkpoint_config is not None:
        # save mmdet version, config file content and class names in
        # checkpoints as meta data
        cfg.checkpoint_config.meta = dict(
            mmdet_version=mmdet_version,
            mmseg_version=mmseg_version,
            mmdet3d_version=mmdet3d_version,
            config=cfg.pretty_text,
            CLASSES=classes,
            PALETTE=datasets[0].PALETTE  # for segmentors
            if hasattr(datasets[0], 'PALETTE') else None)
    # add an attribute for visualization convenience
    model.CLASSES = classes
    # ===== End of Checkpoint Configuration =====

Training Entry Point

The training call is unchanged except that the dataset list may contain a DALI iterator. This case needs to be handled correctly in the training function (see the next section for details).

demos/streampetr/tools/train.py

    # ===== Train the model =====
    #@NOTE
    # Call the training function. Note that datasets is passed to the model. As this now may be a DALI 
    # iterator, one needs to make sure that the training function is able to handle this (by using it directly
    # as a drop-in replacement for a PyTorch DataLoader, instead of creating a PyTorch DataLoader from it). 
    # This is done in `projects/mmdet3d_plugin/core/apis/mmdet_train.py`.
    custom_train_model(
        model,
        datasets,
        cfg,
        distributed=distributed,
        validate=(not args.no_validate),
        timestamp=timestamp,
        meta=meta)
    # ===== End of Train the model =====

Adjustments to the Training Function

In the training script, we call custom_train_model() to train the model, which in turn calls custom_train_detector() (with the latter defined in projects/mmdet3d_plugin/core/apis/mmdet_train.py). The implementation of custom_train_detector() needs to be adjusted to handle the case that the passed dataset is actually a DALIStructuredOutputIterator object. Please see the code snippet below for details on how this is done.

demos/streampetr/projects/mmdet3d_plugin/core/apis/mmdet_train.py

    # ===== Using the DALI Iterator Directly (if used) =====
    #@NOTE 
    # If the passed dataset is a `DALIStructuredOutputIterator` object, use it directly instead of creating a 
    # PyTorch DataLoader from it.
    #
    # Note that the `else` branch is the original implementation. This means that if a PyTorch Dataset is
    # input, the behavior of the original implementation is preserved.
    if isinstance(dataset[0], DALIStructuredOutputIterator):
        data_loaders = dataset
    else:
        data_loaders = [
            build_dataloader(
                ds,
                cfg.data.samples_per_gpu,
                cfg.data.workers_per_gpu,
                # cfg.gpus will be ignored if distributed
                len(cfg.gpu_ids),
                dist=distributed,
                seed=cfg.seed,
                shuffler_sampler=cfg.data.shuffler_sampler,  # dict(type='DistributedGroupSampler'),
                nonshuffler_sampler=cfg.data.nonshuffler_sampler,  # dict(type='DistributedSampler'),
                runner_type=cfg.runner,
                # Use unsharded mini-dataset sampling for the evaluation. This does not represent the real training use-case,
                # but enables the use of the mini-dataset for performance evaluations.
                replicate_dataset_per_gpu=cfg.get('data_loading_config', {}).get(
                    'replicate_dataset_per_gpu', False),
            ) for ds in dataset
        ]
    # ===== End of Using the DALI Iterator Directly (if used) =====

Important

Note that from this point on, the DALIStructuredOutputIterator acts as a drop-in replacement for the original PyTorch DataLoader (in case DALI is used), and no changes deeper in the training code are needed to switch from the original PyTorch DataLoader to the DALI pipeline.

Note

Mini-dataset replication. The performance configurations set data_loading_config.replicate_dataset_per_gpu=True for both loaders because NuScenes mini has too few sequences for the global batch size. This lets each GPU sample the full dataset. It is an evaluation setting, not a required step when integrating DALI into another project; the option defaults to False.

In tools/train.py, get_dali_pipeline_distributed_args() passes rank=0 and world_size=1 to the pipeline setup, while keeping device_id=local_rank. DDP continues to use the real distributed ranks.

For PyTorch, core/apis/mmdet_train.py forwards the flag to datasets/builder.py, which passes rank=0 and num_replicas=1 to the training sequence sampler. Worker initialization still uses the real rank, and validation retains its normal sampler. See the evaluation’s replication policy for the workload and augmentation-seed limitations.