Skip to main content
`Ctrl`+`K`

# Using NHWC Batch Normalization with PyTorch

This topic explains how to use NHWC Batch Normalization for MIOpen operations in PyTorch. NHWC is a deep-learning memory format that has certain performance advantages over traditional memory formats.

For information about installing and using PyTorch with ROCm, see PyTorch on ROCm. For a list of the ROCm components and features that PyTorch supports, see PyTorch compatibility. For more background on using PyTorch and ROCm for AI tasks, see Training a model with PyTorch for ROCm.

## NHWC versus NCHW

NHWC (also known as “Channels Last”) and NCHW are two types of memory formats for deep learning. They describe how multidimensional arrays (nD) are translated to a linear (one-dimensional) memory address space.

- NCHW (Number of samples, channels, height, width): This is the default data layout in which channels are stored separately from one another. The height and width information is stored after the channels.

- NHWC (Number of samples, height, width, channels): In this alternative format, channels are stored next to each other after the height and width information.

The performance of NHWC is better than that of NCHW and is close to that observed when using a blocked memory format. NHWC is also easier to work with for common operations.

For more information about these memory formats, see the PyTorch memory format documentation and the Intel Extension for PyTorch GitHub.

## Batch Normalization

Batch Normalization (also known as Batchnorm or BatchNorm) enables higher learning rates and reduces initialization overhead by normalizing layer inputs. Ordinarily, the distribution of the inputs to each layer changes as the parameters to the previous layer change. This makes it more difficult to train deep learning models and leads to lower learning rates. With Batch Normalization, normalization is part of the architecture and is performed for each training batch.

For more information on Batch Normalization, see Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift.

## Enabling or disabling NHWC Batch Normalization for MIOpen using PyTorch

The PyTorch open-source tensor library provides support for using NHWC Batch Normalization with MIOpen. In addition to Batch Normalization, NHWC support is also available for convolution and other MIOpen features.

NHWC Batch Normalization support in MIOpen can be used in a PyTorch environment using ROCm 7.0 or later. This configuration supports 2D and 3D NHWC Batch Normalization. 1D Batch Normalization is not applicable to the NHWC format.

### PyTorch branch support

The `ROCm/pytorch` PyTorch images support NHWC Batch Normalization. ROCm 7.0 or later is required. The following PyTorch branches support the NHWC Batch Normalization feature:

- release/2.6

- release/2.7

In the `release/2.7` PyTorch branch, NHWC Batch Normalization support in MIOpen is enabled by default. To use the native Batch Normalization approach with this image, use this command:

    PYTORCH_MIOPEN_SUGGEST_NHWC_BATCHNORM=0

In the `release/2.6` PyTorch branch, NHWC Batch Normalization support in MIOpen is disabled by default. To enable NHWC Batch Normalization for this image, use this command:

    PYTORCH_MIOPEN_SUGGEST_NHWC_BATCHNORM=1

For information about installing and using PyTorch on ROCm, see PyTorch on ROCm.

## Supported configurations

The following table shows the Batch Normalization support for NHWC and NCHW with various data types and modes. It also indicates which backend is used with and without the `PYTORCH_MIOPEN_SUGGEST_NHWC_BATCHNORM` environment variable enabled.

Note

Mixed mode means that the Batch Normalization module has a different type than the inputs, for example, an input or gradient data type of `FP16` or `BF16` and a Batch Normalization type of `FP32`. If the Batch Normalization module has the same data type as the inputs, for instance, an input or gradient data type of `FP32` and a Batch Normalization module that is also `FP32`, the mode is “not mixed”.

Input data type | Memory format | Mode | Mixed/not mixed | Backend with NHWC Batch Normalization enabled | Backend with NHWC Batch Normalization disabled
--- | --- | --- | --- | --- | ---
`float32` | NCHW | 1D/2D/3D | not mixed | MIOpen | MIOpen
`float32` | NHWC | 2D/3D | not mixed | native | MIOpen
`float16` | NCHW | 1D/2D/3D | mixed | MIOpen | MIOpen
`float16` | NCHW | 1D/2D/3D | not mixed | native | native
`float16` | NHWC | 2D/3D | mixed | native | MIOpen
`float16` | NHWC | 2D/3D | not mixed | native | native
`bfloat16` | NCHW | 1D/2D/3D | mixed | MIOpen (*) | MIOpen (*)
`bfloat16` | NCHW | 1D/2D/3D | not mixed | native | native
`bfloat16` | NHWC | 2D/3D | mixed | native | MIOpen
`bfloat16` | NHWC | 2D/3D | not mixed | native | native

(*) MIOpen is used with ROCm 6.4 and later. Otherwise, the native backend is used.

## Disabling MIOpen for Batch Normalization in PyTorch

In some situations, you might not want to use MIOpen as the backend for Batch Normalization operations. To disable the use of MIOpen with Batch Normalization, add this code to your application.

    inp = torch.randn(size, requires_grad=True)
    grad = torch.randn(size, requires_grad=False)
    mod = nn.BatchNorm2d(inp.size(1), device="cuda")

    with torch.backends.cudnn.flags(enabled=False): # this line disables MIOpen for the two lines below, native batchnorm will be used

       out = mod(inp)
       out.backward(grad)

## Verifying NHWC Batch Normalization use with MIOpen

For some operations, it can be difficult to determine the backend and memory format used. To verify whether MIOpen is being used and whether the memory format is NHWC or NCHW, run your program with the following environment variable:

    MIOPEN_ENABLE_LOGGING_CMD=1

Here is an example command:

    MIOPEN_ENABLE_LOGGING_CMD=1 python test_nn.py -v -k test_batchnorm_cudnn_nhwc

The output might look like this:

    MIOpen(HIP): Command [LogCmdBNorm] ./bin/MIOpenDriver bnorm -n 4 -c 8 -H 2 -W 2 -m 1 --forw 1 -b 0 -r 1 -s 1 --layout NHWC
    MIOpen(HIP): Command [LogCmdBNorm] ./bin/MIOpenDriver bnorm -n 4 -c 8 -H 2 -W 2 -m 1 --forw 0 -b 1 -s 1 --layout NHWC
    MIOpen(HIP): Command [LogCmdBNorm] ./bin/MIOpenDriver bnorm -n 4 -c 8 -H 2 -W 2 -m 1 --forw 1 -b 0 -r 1 -s 1 --layout NCHW
    MIOpen(HIP): Command [LogCmdBNorm] ./bin/MIOpenDriver bnorm -n 4 -c 8 -H 2 -W 2 -m 1 --forw 0 -b 1 -s 1 --layout NCHW
    MIOpen(HIP): Command [LogCmdBNorm] ./bin/MIOpenDriver bnorm -n 2 -c 8 -H 8 -W 1 -m 1 --forw 1 -b 0 -r 1 -s 1 --layout NHWC
    MIOpen(HIP): Command [LogCmdBNorm] ./bin/MIOpenDriver bnorm -n 2 -c 8 -H 8 -W 1 -m 1 --forw 0 -b 1 -s 1 --layout NHWC
    MIOpen(HIP): Command [LogCmdBNorm] ./bin/MIOpenDriver bnorm -n 2 -c 8 -H 8 -W 1 -m 1 --forw 1 -b 0 -r 1 -s 1 --layout NCHW
    MIOpen(HIP): Command [LogCmdBNorm] ./bin/MIOpenDriver bnorm -n 2 -c 8 -H 8 -W 1 -m 1 --forw 0 -b 1 -s 1 --layout NCHW

Each line corresponds to a different command or operation. The `./bin/MIOpenDriver` string indicates that MIOpen was used for the operation. The `--layout` parameter shows whether NHWC or NCHW was used, for example, `--layout NHWC` means the NHWC memory format was used.

## Running Batch Normalization tests

Several test suites are available for Batch Normalization. To test Batch Normalization training using both NHWC and NCHW in 2D, run the following command:

    python test_nn.py -v -k test_batchnorm_2D_train

To test Batch Normalization training using both NHWC and NCHW in 3D, run the following command:

    python test_nn.py -v -k test_batchnorm_3D_train

To test Batch Normalization inference for 2D using both memory formats, use this command:

    python test_nn.py -v -k test_batchnorm_2D_inference

To test the same functionality for 3D, use this command:

    python test_nn.py -v -k test_batchnorm_3D_inference