How to Convert Between NHWC and NCHW in TensorFlow

As a data scientist, you may often encounter the need to convert between different data formats when working with machine learning models. One common conversion is between the NHWC and NCHW formats in TensorFlow. In this article, we will explain what these formats are, why they are important, and how to perform the conversion in TensorFlow.

Table of Contents

1. What are NHWC and NCHW?
2. Why is Conversion Important?
3. How to Convert Between NHWC and NCHW in TensorFlow
4. Common Errors and How to Handle Them
5. Conclusion

What are NHWC and NCHW?

Before diving into the conversion methods, let’s briefly review NHWC and NCHW formats.

• NHWC (TensorFlow Default):

  • N: Number of samples
  • H: Height of the input data
  • W: Width of the input data
  • C: Number of channels (e.g., RGB for images)

• NCHW:

  • N: Number of samples
  • C: Number of channels
  • H: Height of the input data
  • W: Width of the input data

The NHWC format is the default format used by TensorFlow and many other deep learning frameworks. In this format, the tensor is represented as a four-dimensional array with dimensions (batch_size, height, width, channels). The batch_size dimension represents the number of examples in the input data, while the height, width, and channels dimensions represent the spatial dimensions and number of channels in the data.

The NCHW format, on the other hand, is commonly used by deep learning frameworks such as Caffe and PyTorch. In this format, the tensor is represented as a four-dimensional array with dimensions (batch_size, channels, height, width). The batch_size dimension represents the number of examples, while the channels dimension represents the number of channels in the data. The height and width dimensions represent the spatial dimensions of the data.

Why is Conversion Important?

Conversion between NHWC and NCHW is important when working with deep learning models that use different formats for their input and output data. For example, if you are using a pre-trained model that uses the NCHW format, but your input data is in the NHWC format, you will need to convert the data before passing it through the model.

How to Convert Between NHWC and NCHW in TensorFlow

TensorFlow provides two functions to convert between NHWC and NCHW formats: tf.transpose and tf.compat.v1.nn.conv2d_transpose.

Using tf.transpose

The tf.transpose function can be used to permute the dimensions of a tensor. To convert from NHWC to NCHW, we need to permute the dimensions from (batch_size, height, width, channels) to (batch_size, channels, height, width). Similarly, to convert from NCHW to NHWC, we need to permute the dimensions from (batch_size, channels, height, width) to (batch_size, height, width, channels).

Here is an example of how to use tf.transpose to convert a tensor from NHWC to NCHW:

import tensorflow as tf

# input tensor in NHWC format
input_nhwc = tf.random.uniform([10, 28, 28, 3])
print("Shape before conversion: \n", input_nhwc.shape)
# convert to NCHW format
input_nchw = tf.transpose(input_nhwc, perm=[0, 3, 1, 2])
print("Shape after conversion: \n", input_nchw.shape)

Output:

Shape before conversion: 
 (10, 28, 28, 3)
Shape after conversion: 
 (10, 3, 28, 28)

In this example, we first create a random tensor in NHWC format with shape (batch_size, height, width, channels). We then use tf.transpose to permute the dimensions according to the perm argument, which specifies the new order of dimensions. The resulting tensor, input_nchw, has shape (batch_size, channels, height, width).

Similarly, here is an example of how to use tf.transpose to convert a tensor from NCHW to NHWC:

import tensorflow as tf

# input tensor in NHWC format
input_nhwc = tf.random.uniform([10, 3, 28, 28])
print("Shape before conversion: \n", input_nhwc.shape)
# convert to NCHW format
input_nchw = tf.transpose(input_nhwc, perm=[0, 2, 3, 1])
print("Shape after conversion: \n", input_nchw.shape)

Output:

Shape before conversion: 
 (10, 3, 28, 28)
Shape after conversion: 
 (10, 28, 28, 3)

In this example, we create a new random tensor in NCHW format with shape (batch_size, channels, height, width). We then use tf.transpose to permute the dimensions according to the perm argument, which specifies the new order of dimensions. The resulting tensor, input_nhwc, has shape (batch_size, height, width, channels).

Using tf.reshape

The tf.reshape function allows you to change the shape of a tensor while preserving its elements. For NHWC to NCHW conversion, you can use this function to rearrange the dimensions.

Here is an example of how to use tf.reshape to convert a tensor from NHWC to NCHW:

import tensorflow as tf

# input tensor in NHWC format
input_nhwc = tf.random.uniform([10, 28, 28, 3])
# get shape
shape = tf.shape(input_nhwc)
print("Shape before conversion: \n", input_nhwc.shape)
# convert to NCHW format
input_nchw = tf.reshape(input_nhwc, [shape[0], shape[3], shape[1], shape[2]])
print("Shape after conversion: \n", input_nchw.shape)

Output:

Shape before conversion: 
 (10, 28, 28, 3)
Shape after conversion: 
 (10, 3, 28, 28)

In this example, we first create a random input tensor in NHWC format with shape (batch_size, height, width, channels). We then use tf.reshape to perform a deconvolution operation on the input tensor. The resulting tensor, input_nchw, has shape (batch_size, channels, height, width). Converting a tensor from NCHW to NHWC can be done in the same manner.

Common Errors and How to Handle Them:

Shape Mismatch:

• Error: "Shape must be rank 4 but is rank X."
• Solution: Ensure that the input tensor has the correct rank (4 for NHWC and NCHW).

Incorrect Permutation:

• Error: "Invalid permutation."
• Solution: Double-check the order of dimensions in the perm argument to tf.transpose.

Memory Issues:

• Error: "OOM when allocating tensor with shape."
• Solution: Be mindful of the available memory when working with large tensors, and consider batching or using smaller models.

Conclusion

In this article, we explained what the NHWC and NCHW formats are, why they are important, and how to convert between them in TensorFlow. We showed two methods for performing the conversion: using tf.transpose to permute the dimensions of a tensor, and using tf.compat.v1.nn.conv2d_transpose to perform a deconvolution operation. By understanding these methods, you can easily convert between NHWC and NCHW formats and work with deep learning models that use different formats for their input and output data.

SHARE:

About Saturn Cloud

Saturn Cloud is your all-in-one solution for data science & ML development, deployment, and data pipelines in the cloud. Spin up a notebook with 4TB of RAM, add a GPU, connect to a distributed cluster of workers, and more. Request a demo today to learn more.