[Torchvision API] Support for volumetric data by mdabek-nvidia · Pull Request #6544 · NVIDIA/DALI

mdabek-nvidia · 2026-10-22T12:35:35Z

Category:

Bug fix

Description:

Torchvision.v2 accepts arbitrary leading dimensions: (..., C,H,W) the current implementation of Torchvision API uses only CHW or NCHW. That does not allow to process other data layouts like NFCHW (e.g. video frames)

The default behavior, to match the original Torchvision is to set CHW or FCHW layouts for tensors, depending on the first input data (changing layout mid-pipeline run would cause huge performance penalty, so we are not doing that). Next, the API is extended (by explicit layout and batch size arguments) to allow passing NCHW and NFCHW.

Currently, we do infer HWC only from PIL images and they do not support batching, so we do not allow: FHWC, NHWC or NFHWC, since it is unlikely to find them in the existing SW, that currently uses Torchvision. However, it is possible that these layouts will be implemented in the future as the enhancement to the existing Torchvison capabilities.

Additional information:

Affected modules and functionalities:

Reshaping the input and recovering the input shape were added to the data layout transformation.

Key points relevant for the review:

Tests:

Checklist

Documentation

DALI team only

Requirements

Implements new requirements
Affects existing requirements
N/A

REQ IDs: N/A

JIRA TASK: DALI-4891

Signed-off-by: Marek Dabek <mdabek@nvidia.com>

copy-pr-bot · 2026-10-22T12:35:39Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Signed-off-by: Marek Dabek <mdabek@nvidia.com>

review-notebook-app · 2026-10-22T12:36:25Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

mdabek-nvidia · 2026-10-22T12:38:47Z

!build

greptile-apps · 2026-10-22T12:40:31Z

Confidence Score: 5/5

[Medium impact] The latest changes appear safe to merge, with a non-blocking mismatch between the documented and actual unsupported-layout exception.

Findings

Unsupported layouts raise another error ▶

Summary

Adds higher-rank tensor support to the TorchVision object API.

Torchvision transforms keep frame sequences and batches in their original shape.
Normalize works across frame layouts and rejects PIL images.

The unsupported-layout exception no longer matches the guide. mdabek-nvidia rejected adding layout enums because DALI has no existing layout enums; Greptile withdrew that finding.

_{Reviews (8) · Last reviewed commit: "Review fixeS" · Reviewed by Greptile}

dali-automaton · 2026-10-22T12:43:00Z

CI MESSAGE: [75843191]: BUILD STARTED

TorchVision samples a single parameter set per call and applies it across all leading dimensions (..., C, H, W). Compose drew random values per sample instead, so images in a batch and frames of a video were transformed differently. PipelineCHW now flattens the leading dimensions into frames of FCHW sequences: - GPU: the whole call is processed as a single sample, which is also faster than per-sample processing (1.2-1.9x in RRC+flip+normalize). - CPU: frames are split into up to min(batch_size, num_threads) sequences to keep thread parallelism. Random values drawn for the first sequence are broadcast with fn.permute_batch through the new share_random_value() helper. RandomResizedCrop on CPU generates the crop window once with fn.random_crop_generator and resizes with an ROI. Normalize uses a 4D mean/std for FCHW input. batch_size no longer limits the leading dimensions of the input. Signed-off-by: Marek Dabek <mdabek@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

dali-automaton · 2026-10-22T14:44:37Z

CI MESSAGE: [75843191]: BUILD FAILED

dali-automaton · 2026-10-23T07:55:51Z

CI MESSAGE: [75843191]: BUILD PASSED

Signed-off-by: Marek Dabek <mdabek@nvidia.com>

mdabek-nvidia · 2026-10-23T13:58:12Z

!build

dali-automaton · 2026-10-23T14:03:03Z

CI MESSAGE: [76146454]: BUILD STARTED

mdabek-nvidia · 2026-10-23T14:27:01Z

!build

dali-automaton · 2026-10-23T14:33:06Z

CI MESSAGE: [76151558]: BUILD STARTED

dali-automaton · 2026-10-23T16:15:53Z

CI MESSAGE: [76151558]: BUILD PASSED

dali-automaton · 2026-10-23T16:43:25Z

CI MESSAGE: [76146454]: BUILD PASSED

This reverts commit f1e65fc.

This reverts commit 261dc7a.

This reverts commit 89af03c.

* Supporting CHW and FCHW as default layouts for tensors * Supporting NCHW and NFCHW as extensions to default layouts Signed-off-by: Marek Dabek <mdabek@nvidia.com>

mdabek-nvidia · 2026-10-26T15:26:03Z

@greptile-apps please review

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

Signed-off-by: Marek Dabek <mdabek@nvidia.com>

mdabek-nvidia · 2026-10-27T12:49:12Z

@greptile-apps please re-review

mdabek-nvidia · 2026-10-27T12:49:30Z

!build

dali-automaton · 2026-10-27T12:58:02Z

CI MESSAGE: [77011372]: BUILD STARTED

Signed-off-by: Marek Dabek <mdabek@nvidia.com>

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

mdabek-nvidia · 2026-10-28T09:43:10Z

!build

dali-automaton · 2026-10-28T09:50:37Z

CI MESSAGE: [77264341]: BUILD STARTED

mzient · 2026-10-28T14:49:17Z

+        elif self.sample_layout == "FCHW":
+            # Any leading dimensions are flattened into the frames dimension, as in torchvision
+            tail = 3
+            supported = data_input.ndim > tail
+            expected = "at least 4 dimensions"
        else:
-            raise ValueError("CHW layout is currently supported for torch.Tensor only.\
-                Please check if samples have the same format.")
+            tail = 3
+            supported = data_input.ndim == tail
+            expected = "exactly 3 dimensions"
+        if not supported:


Suggested change

elif self.sample_layout == "FCHW":

# Any leading dimensions are flattened into the frames dimension, as in torchvision

tail = 3

supported = data_input.ndim > tail

expected = "at least 4 dimensions"

else:

raise ValueError("CHW layout is currently supported for torch.Tensor only.\

Please check if samples have the same format.")

tail = 3

supported = data_input.ndim == tail

expected = "exactly 3 dimensions"

if not supported:

else:

tail = 3 # we always need at least CHW

# Any leading dimensions are flattened into the frames dimension, as in torchvision

expected_dims = 4 if sample_layout == "FCHW" else 3

supported = data_input.ndim >= expected_dims

expected = f"at least {expected_dims} dimensions"

if not supported:

That's actually invalid fix - the supported will pass layout "CHW" and input tensor (4, 3, 8, 8,) leading to DALI throwing exception about the invalid batch size.
I've managed to simplify the code anyway.

mzient · 2026-10-28T15:39:39Z

+    "NFCHW": (PipelineFCHW, True),
+}
+_UNSUPPORTED_LAYOUTS = ("NHWC", "FHWC", "NFHWC")
+_ALLOWED_LAYOUTS = (None, "HWC", "CHW", "NHWC", "NCHW", "FCHW", "FHWC", "NFCHW", "NFHWC")


As a user I wouldn't be happy if the program first showed me an "allowed" layout only to find out it's not supported.

That's correct - I removed unsupported layouts from allowed layouts, but I am leaving the _UNSUPPORTED_LAYOUTS, for the future support.

mzient · 2026-10-28T15:41:21Z

        **dali_pipeline_kwargs,
    ):
+        if layout not in _ALLOWED_LAYOUTS:
+            raise ValueError(f"Unknown layout {layout!r}. Allowed values: {_ALLOWED_LAYOUTS}.")


First we say a layout is allowed, but it may yet be unsupported. This is user-hostile.

Changed the _ALLOWED_LAYOUTS to reflect what is supported.

Signed-off-by: Marek Dabek <mdabek@nvidia.com>

greptile-apps · 2026-10-28T18:42:18Z

+    "NFCHW": (PipelineFCHW, True),
+}
+_UNSUPPORTED_LAYOUTS = ("NHWC", "FHWC", "NFHWC")
+_ALLOWED_LAYOUTS = (None, "HWC", "CHW", "NCHW", "FCHW", "NFCHW")


Unsupported layouts raise another error

[Bug] Removing NHWC, FHWC, and NFHWC from _ALLOWED_LAYOUTS makes them raise ValueError before reaching the NotImplementedError check. The input-layout guide still promises NotImplementedError, so users following that guide would catch the wrong exception. Keep the documented exception or update the guide and remove the unreachable check.

Suggested change

_ALLOWED_LAYOUTS = (None, "HWC", "CHW", "NCHW", "FCHW", "NFCHW")

_ALLOWED_LAYOUTS = (None, "HWC", "CHW", "NHWC", "NCHW", "FCHW", "FHWC", "NFCHW", "NFHWC")

dali-automaton · 2026-10-28T22:19:22Z

CI MESSAGE: [77264341]: BUILD FAILED

dali-automaton · 2026-10-29T11:35:28Z

CI MESSAGE: [77264341]: BUILD PASSED

Support for volumetric data

923f516

Signed-off-by: Marek Dabek <mdabek@nvidia.com>

Documentation update

f9c60f3

Signed-off-by: Marek Dabek <mdabek@nvidia.com>

greptile-apps Bot reviewed Oct 22, 2026

View reviewed changes

Comment thread dali/test/python/torchvision/test_tv_compose.py Outdated

mdabek-nvidia mentioned this pull request Oct 23, 2026

[Torchvision] Torchvision per sample random in batches #6548

Draft

4 of 18 tasks

Review fixed in tests

261dc7a

Signed-off-by: Marek Dabek <mdabek@nvidia.com>

greptile-apps Bot reviewed Oct 23, 2026

View reviewed changes

dali-automaton assigned banasraf and mzient Oct 23, 2026

Review fixes

89af03c

Signed-off-by: Marek Dabek <mdabek@nvidia.com>

mdabek-nvidia marked this pull request as draft October 26, 2026 08:38

mdabek-nvidia added 4 commits October 26, 2026 14:39

Revert "Sample one random parameter set per call in torchvision Compose"

4c0b30d

This reverts commit f1e65fc.

Revert "Review fixed in tests"

d40c7de

This reverts commit 261dc7a.

Revert "Review fixes"

1ce7101

This reverts commit 89af03c.

Cleaned up layouts support

a1c5011

* Supporting CHW and FCHW as default layouts for tensors * Supporting NCHW and NFCHW as extensions to default layouts Signed-off-by: Marek Dabek <mdabek@nvidia.com>

greptile-apps Bot reviewed Oct 26, 2026

View reviewed changes

Comment thread dali/python/nvidia/dali/experimental/torchvision/v2/normalize.py

Comment thread dali/test/python/torchvision/test_tv_compose.py Outdated

Comment thread dali/python/nvidia/dali/experimental/torchvision/v2/compose.py

mdabek-nvidia and others added 2 commits October 27, 2026 13:39

Enforcing strict length in zip

96f00d8

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

Black formatting

4ef22dd

Signed-off-by: Marek Dabek <mdabek@nvidia.com>

Documentation and normalize error for PIl images

9059141

Signed-off-by: Marek Dabek <mdabek@nvidia.com>

mdabek-nvidia marked this pull request as ready for review October 27, 2026 19:15

greptile-apps Bot reviewed Oct 27, 2026

View reviewed changes

Comment thread dali/test/python/torchvision/test_tv_normalize.py Outdated

Assert message checking

80737a2

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

mzient reviewed Oct 28, 2026

View reviewed changes

Review fixeS

f98758b

Signed-off-by: Marek Dabek <mdabek@nvidia.com>

greptile-apps Bot reviewed Oct 28, 2026

View reviewed changes

	_ALLOWED_LAYOUTS = (None, "HWC", "CHW", "NCHW", "FCHW", "NFCHW")
	_ALLOWED_LAYOUTS = (None, "HWC", "CHW", "NHWC", "NCHW", "FCHW", "FHWC", "NFCHW", "NFHWC")

Conversation

mdabek-nvidia commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Category:

Description:

Additional information:

Affected modules and functionalities:

Key points relevant for the review:

Tests:

Checklist

Documentation

DALI team only

Requirements

Uh oh!

copy-pr-bot Bot commented Oct 22, 2026

Uh oh!

review-notebook-app Bot commented Oct 22, 2026

Uh oh!

mdabek-nvidia commented Oct 22, 2026

Uh oh!

greptile-apps Bot commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Findings

Summary

Uh oh!

Uh oh!

dali-automaton commented Oct 22, 2026

Uh oh!

dali-automaton commented Oct 22, 2026

Uh oh!

dali-automaton commented Oct 23, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mdabek-nvidia commented Oct 23, 2026

Uh oh!

dali-automaton commented Oct 23, 2026

Uh oh!

mdabek-nvidia commented Oct 23, 2026

Uh oh!

dali-automaton commented Oct 23, 2026

Uh oh!

dali-automaton commented Oct 23, 2026

Uh oh!

dali-automaton commented Oct 23, 2026

Uh oh!

mdabek-nvidia commented Oct 26, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mdabek-nvidia commented Oct 27, 2026

Uh oh!

mdabek-nvidia commented Oct 27, 2026

Uh oh!

dali-automaton commented Oct 27, 2026

Uh oh!

Uh oh!

mdabek-nvidia commented Oct 28, 2026

Uh oh!

dali-automaton commented Oct 28, 2026

Uh oh!

mzient Oct 28, 2026

Choose a reason for hiding this comment

Uh oh!

mdabek-nvidia Oct 28, 2026

Choose a reason for hiding this comment

Uh oh!

mzient Oct 28, 2026

Choose a reason for hiding this comment

Uh oh!

mdabek-nvidia Oct 28, 2026

Choose a reason for hiding this comment

Uh oh!

mzient Oct 28, 2026

Choose a reason for hiding this comment

Uh oh!

mdabek-nvidia commented Oct 22, 2026 •

edited

Loading

greptile-apps Bot commented Oct 22, 2026 •

edited

Loading