Skip to content

torch conv does two NHWC<->NCHW copies when one would do #23279

Description

@pctablet505

On the torch backend, Conv (keras/src/backend/torch/nn.py, conv()) makes two input copies per forward pass in the default channels_last data format: _transpose_spatial_inputs permutes NHWC->NCHW and calls a plain .contiguous(), then _maybe_convert_to_channels_last immediately calls .contiguous(memory_format=torch.channels_last) on the result.

The first copy is redundant every single call: a permuted NHWC->NCHW view of a contiguous input is already channels_last-contiguous, so the standard-contiguous tensor produced by step one is never consumed -- the very next line reconverts it to channels_last anyway. This isn't data-dependent; it's a fixed property of the transpose, so it fires on every conv forward regardless of input.

Fix: collapse the two steps into one -- have _transpose_spatial_inputs permute to a view and call .contiguous(memory_format=torch.channels_last / channels_last_3d) directly for the 4-D/5-D conv path, and drop the now-redundant separate _maybe_convert_to_channels_last call in conv(). Makes the common case zero-copy on input.

Part of the per-call Python-dispatch overhead series in #22561. Fixed by #23199.

Metadata

Metadata

Assignees

Labels

Type

No type

Fields

No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions