Channels Last (NHWC) support plan.

[VitalyFedyunin]

Release goals:

• Support ResNet50, AlexNet, ResNeXt and proof performance gain on mixed precision models + Volta GPUs.
• Unlock developers to extend channels last operators coverage.

Tasks:

• Change factory functions (_like and similars) to optionally preserve format (Fix jit outplace tracing and reapply changes to *_like operators. #28839) and land with docs.
• Switch factory functions to preserve memory format by default (BC breaking) and fix all call sites (link tracking issue).
• Fix TensorIterator to preserve memory format (TensroIterator preserve format for binary, ternary operators. #28291).
• Land CudNN changes to support channels last ([cudnn nhwc support] #23861).
• Land adaptive_avg_pool2d (CUDA) changes to support channels last ([nhwc support for adaptive_avg_pool2d & adaptive_avg_pool2d_backward] #24396).
• Add torch.memory_format support to the TorchScript (Add torch.memory_format support to the TorchScript #28544).
• Add model developer document how to check if model propagates NHWC forward and backward. Add documentation how to debug if model loosing channels last.
• Add TorchVision transformer and loader PRs (need to avoid additional memory copies on loading).
• Add developer docs how to write NHWC operators (https://github.com/pytorch/pytorch/wiki/Writing-memory-format-aware-operators)
• Add dev docs how to check if operator support memory formats.
• Documentation describing what is channels last, overlapping tensors, dense tensors, contiguous and when tensor can be both.
• Add channels last support to cuda.comm.scatter and cuda.comm.gather (Add channels last support to cuda.comm.scatter and gather #28077)
• Add model level .to_channels_last operator to convert weights.
• Fix max_pool2d bug (Channels last max_pool2d fails with specific inputs #28714).
• Revoke mutually exclusive requirement on channels last and contiguous memory formats (Revoking mutually exclusive requirement on channels last and contiguous tensor #28466).
• Figure out if we can ship binaries with CudNN >= 7.4
• Update site docs for *_like, clone, to

Stretch goals:

• 1d and 3d support.

cc @VitalyFedyunin

[thuyen]

@VitalyFedyunin will this be part of 1.4?

[VitalyFedyunin]

Unfortunately not, as we met few problems with CudNN and NC11 tensors, but expect it in master first month of 2020.

[thuyen]

Got it. Thanks!

[ketjow4]

Hey, @VitalyFedyunin what is current status? What are your plans?

[VitalyFedyunin]

It is done in master and you can try using channels last for mixed precision training (for example NVIDIA/apex#668). Documentation updates will follow soon.

[ketjow4]

Thanks for update @VitalyFedyunin. Will I be able to train model using single precision with NHWC data layout ? I am asking as I search for possibility of converting pytroch model through onnx to tensorflow. Is there any list of supported operators for NHWC training?

[codeislife99]

Hi, Which pytorch version will this update be a part of ?

[VitalyFedyunin]

It is out with 1.5 you can read more here https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html