Skip to content

Commit

Permalink
[docs] fix incorrect example in convert_conv3d_weight_memory_format (
Browse files Browse the repository at this point in the history
…#129318)

The current example fails when using `torch.channels_last`, and the docs are slightly incorrect for the 3d case.
Pull Request resolved: #129318
Approved by: https://github.com/albanD
  • Loading branch information
a-r-r-o-w authored and pytorchmergebot committed Aug 7, 2024
1 parent 6a348e5 commit 525fdc0
Showing 1 changed file with 8 additions and 8 deletions.
16 changes: 8 additions & 8 deletions torch/nn/utils/memory_format.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,13 @@ def convert_conv2d_weight_memory_format(module, memory_format):
layer with 4d weight will be affected by ``model.to``, which does not
necessarily benefit from conversion to specified ``memory_format``.
One place we are confident in is that NHWC(channels_last) conversion for
convolution in cuDNN, As it is beneficial to run convolution in NHWC,
convolution in cuDNN, as it is beneficial to run convolution in NHWC,
even in cases where we have to apply permutation to input tensors.
Hence our strategy here is to convert only the weight of convolution to
channels_last. This ensures that;
1. Fast convolution kernels will be used, the benefit of which could
outweigh overhead of permutation (if input is not in the same format)
outweigh overhead of permutation (if input is not in the same format).
2. No unnecessary permutations are applied on layers that do not benefit
from memory_format conversion.
Expand Down Expand Up @@ -85,18 +85,18 @@ def convert_conv3d_weight_memory_format(module, memory_format):
provides considerable speed up for fp16 data on CUDA devices with compute capability >= 7.0
.. note::
Calling ``model.to(memory_format=torch.channels_last)`` is more aggressive
Calling ``model.to(memory_format=torch.channels_last_3d)`` is more aggressive
than the utility function ``convert_conv3d_weight_memory_format``. Any
layer with 4d weight will be affected by ``model.to``, which does not
necessarily benefit from conversion to specified ``memory_format``.
One place we are confident in is that NHWC(channels_last) conversion for
convolution in cuDNN, As it is beneficial to run convolution in NHWC,
One place we are confident in is that NDHWC(channels_last_3d) conversion for
convolution in cuDNN, as it is beneficial to run convolution in NDHWC,
even in cases where we have to apply permutation to input tensors.
Hence our strategy here is to convert only the weight of convolution to
channels_last. This ensures that;
channels_last_3d. This ensures that;
1. Fast convolution kernels will be used, the benefit of which could
outweigh overhead of permutation (if input is not in the same format)
outweigh overhead of permutation (if input is not in the same format).
2. No unnecessary permutations are applied on layers that do not benefit
from memory_format conversion.
Expand Down Expand Up @@ -133,7 +133,7 @@ def convert_conv3d_weight_memory_format(module, memory_format):
>>> model = nn.Sequential(
>>> nn.Conv3d(8, 4, 3)).cuda().half()
>>> # This is identical to:
>>> # nn.utils.convert_conv3d_weight_memory_format(model, torch.channels_last)
>>> # nn.utils.convert_conv3d_weight_memory_format(model, torch.channels_last_3d)
>>> model = nn.utils.convert_conv3d_weight_memory_format(model, torch.channels_last_3d)
>>> out = model(input)
"""
Expand Down

0 comments on commit 525fdc0

Please sign in to comment.