|
|
import torch |
|
|
|
|
|
def convert_conv2d_weight_memory_format(module, memory_format): |
|
|
r"""Convert ``memory_format`` of ``nn.Conv2d.weight`` to ``memory_format`` |
|
|
The conversion recursively applies to nested ``nn.Module``, including ``module``. |
|
|
Note that it only changes the memory_format, but not the semantics of each dimensions. |
|
|
This function is used to facilitate the computation to adopt NHWC kernels, which |
|
|
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 |
|
|
than the utility function ``convert_conv2d_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, |
|
|
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) |
|
|
2. No unnecessary permutations are applied on layers that do not benefit |
|
|
from memory_format conversion. |
|
|
|
|
|
The optimal case is that, layers between convolution layers are channels |
|
|
last compatible. Input tensor would be permuted to channels last when it |
|
|
encounters the first convolution layer and stay in that memory format. |
|
|
Hence following convolutions will not need to permute its input tensor. |
|
|
|
|
|
In case where a channels last incompatible layer is between convolution |
|
|
layers, we need to permute the input tensor back to contiguous format |
|
|
for that layer. The input tensor will go through the remaining layers in |
|
|
contiguous format and be permuted to channels last when it encounters |
|
|
another convolution layer. There's no point in propagating that |
|
|
permutation to an earlier layer, as most layers are quite agnostic to |
|
|
``memory_format``. |
|
|
|
|
|
This claim might change when PyTorch supports fusion of permutation, as |
|
|
there might have been a better spot to fuse the permutation other than |
|
|
immediately before a convolution. |
|
|
|
|
|
Args: |
|
|
module (nn.Module): ``nn.Conv2d`` & ``nn.ConvTranspose2d`` or container |
|
|
``nn.Module`` |
|
|
format: user specified ``memory_format``, |
|
|
e.g. ``torch.channels_last`` or ``torch.contiguous_format`` |
|
|
|
|
|
Returns: |
|
|
The original module with updated ``nn.Conv2d`` |
|
|
|
|
|
Example: |
|
|
>>> # xdoctest: +REQUIRES(env:CUBLAS_WORKSPACE_CONFIG) |
|
|
>>> input = torch.randint(1, 10, (2, 8, 4, 4), dtype=torch.float16, device="cuda") |
|
|
>>> model = nn.Sequential( |
|
|
>>> nn.Conv2d(8, 4, 3)).cuda().half() |
|
|
>>> # This is identical to: |
|
|
>>> # nn.utils.convert_conv2d_weight_memory_format(model, torch.channels_last) |
|
|
>>> model = nn.utils.convert_conv2d_weight_memory_format(model, torch.channels_last) |
|
|
>>> out = model(input) |
|
|
""" |
|
|
|
|
|
|
|
|
|
|
|
if isinstance(module, torch.nn.Conv2d) or isinstance(module, torch.nn.ConvTranspose2d): |
|
|
weight_data = module.weight.detach().clone().contiguous(memory_format=memory_format) |
|
|
module.weight.data = weight_data.resize_(weight_data.size(), memory_format=memory_format) |
|
|
for child in module.children(): |
|
|
convert_conv2d_weight_memory_format(child, memory_format) |
|
|
return module |
|
|
|
