Evaluation

This page summarizes the performance impact of batching the loss computation using accvlab.batching_helpers in a StreamPETR training setup.

See also

For the implementation changes behind these results, see the Batching Helpers demo and its tutorials on batching the matching and loss computations. The Optimization Testing Tools demo explains the runtime measurements and profiling support used in the evaluation.

Setup

Experiment Setup

For this evaluation, the training is performed for the NuScenes mini dataset, using the StreamPETR model. The adaptations to the original StreamPETR implementation are as follows:

StreamPETR – Optimization Overview

  • HungarianAssigner3D & HungarianAssigner2D: Original matchers operate on a per‑sample basis.

    • Cost matrix computation (pre‑requisite for actual matching) → Optimized (following the approach outlined in the batching-helpers example)

    • Matching itself (SciPy implementation on the CPU).

    • HungarianAssigner3D: nan_to_num() was a bottleneck; moved to GPU → Changed in both reference & optimized.

  • StreamPETR Head

    • The reference already concatenates the samples and calls each classification and bounding-box loss function once per decoder layer. The loss functions themselves were therefore already batched over samples.

    • Target preparation around those loss calls is further batched → Optimized. The optimized head keeps the sample dimension in RaggedBatch objects, computes the assignment cost matrices and performs target indexing in batches instead of calling _get_target_single() for each sample, and computes bounding-box loss only for matched pairs. The SciPy Hungarian solver itself still runs once for each layer/sample entry.

    • Loss computation is batched across decoder layers → Optimized. The reference calls loss_single() and dn_loss_single() separately for each layer. The optimized head processes the layer and sample dimensions together in one loss-function call, then reduces over samples and objects while retaining separate loss values for each layer. This also applies to the DN losses.

  • Focal Head

    • Loss computation is already batched over samples & camera images

    • Can use a batched assigner → Optimized.

    • Added use of the custom Gaussian heatmap generator → Changed in both reference & optimized.

Note

Some of the changes/optimizations are done in both the reference and optimized implementation (indicated in the overview above). These changes are not specific to the batching optimization, and are therefore applicable to both implementations. Applying them in both ensures a fair comparison, with the obtained speedup reflecting the effect of the batching optimizations.

See also

For a general overview on how to use the batching helpers to optimize the loss computation, please refer to the example.

The evaluation is performed for a batch size of 8 (in contrast to the original configuration of 2) to obtain a realistic setup and highlight the performance impact of the batching in this case.

Hardware Setup

The hardware setup used for the evaluation is as follows:

System Configuration

GPU

CPU

NVIDIA A100-SXM4-80GB

2x AMD EPYC 7742 64-core Processors

Results

The numbers below are averages over 950 measured iterations after 50 warm-up iterations.

In the following table, the runtime breakdown of the training iteration is shown. Note that the grayed out cells do not contain any optimized code. While theoretically, the Optimization (including Backward) step contains changes (due to the different implementation of the loss, leading to different steps in the backward propagation), the resulting runtime differences in this step are negligible. The Remaining columns in the table show the runtime of the parts of the implementation for which the runtime is not measured directly. For the Forward pass, this mostly corresponds to the forward pass through the network itself (as opposed to the loss computation). For the Training Iteration, this may correspond to additional overhead such as e.g. obtaining/waiting for the next batch of data.

Runtime Baseline → Runtime Optimized [ms] (Speedup ×-fold)
Training Iteration
737 → 602 (× 1.22)
Forward (including Loss) Optimization (including Backward) Remaining
344 → 210 (× 1.64) 311 → 307 81 → 84
Remaining Loss
169 → 170 175 → 40 (× 4.35)
StreamPETR Head Focal Head
79 → 24 (× 3.22) 96 → 15 (× 6.24)

The batching of the loss computation leads to a speedup of × 4.35 for the loss computation itself, with different speedups achieved for the different types of loss. The loss optimization translates to an overall speedup of × 1.64 for the forward pass and × 1.22 for the training iteration. Note that the expected speedup strongly depends on the used batch size.

How to Run

Running the Evaluated Variants

Note

This section assumes that you are running the evaluation in the Docker container set up as described in the StreamPETR Docker setup, with the StreamPETR implementation located at /streampetr. Configuration paths below are relative to demos/streampetr/projects/configs/StreamPETR/.

The evaluation was performed with these two configurations:

Evaluated Batching Helpers configurations

Configuration

Purpose

Evals/batching_helpers_ref.py

Reference implementation with the same GPU nan_to_num and custom heatmap improvements as the optimized configuration.

Evals/batching_helpers_test.py

Batched assigners and loss computation. Comparing it with the reference isolates the batching gains reported above.

Run the configurations individually with:

# Reference implementation with the same unrelated optimizations as the test configuration.
# Comparing test against ref therefore isolates the performance gains from batching.
python /streampetr/tools/train.py --no-validate /streampetr/projects/configs/StreamPETR/Evals/batching_helpers_ref.py

# Implementation with batched loss computation, including the same unrelated optimizations.
python /streampetr/tools/train.py --no-validate /streampetr/projects/configs/StreamPETR/Evals/batching_helpers_test.py

Additional Variant

Evals/batching_helpers_orig_impl.py is a supplemental near-original variant that uses CPU nan_to_num and the original heatmap generation. Comparing the optimized configuration against this variant also includes gains unrelated to batching, so it is not the baseline for the batching speedups reported above.

Run it with:

python /streampetr/tools/train.py --no-validate /streampetr/projects/configs/StreamPETR/Evals/batching_helpers_orig_impl.py

Running All Variants and Storing the Results

The demos/streampetr/tools/run_batching_eval.sh script runs the two evaluated configurations and the supplemental near-original variant. Run it with:

cd /streampetr/tools
bash run_batching_eval.sh

The script stores the results in the res_batching_eval folder in the current working directory. Refer to the script for more details.

Profiling

The demos/streampetr/tools/run_batching_prof.sh script captures Nsight Systems profiles for the implementation with batched loss computation, the reference with GPU nan_to_num, and the reference variant with CPU nan_to_num. It uses the configurations under Profiling/ with 704 × 256 images, 50 training iterations, fixed augmentations and a fixed sample order.

Run it inside the prepared demo container, with the dataset available and the Nsight Systems nsys command on PATH:

cd /streampetr/tools
bash run_batching_prof.sh

The script writes reports to the current working directory with the prefixes streampetr_batched, streampetr_orig_switched (the reference with GPU nan_to_num), and streampetr_orig (the variant with CPU nan_to_num). Open the reports in Nsight Systems to inspect the NVTX ranges. See the profiling guide for details of the instrumentation.