How to convert Tensorflow graph from `NHWC` format to `NCHW`?


#1

I don’t understand how to convert Tensorflow graph from NHWC format to NCHW format to be used in TVM software stack.

  1. Following this tutorial, I use nnvm.frontend.from_tensorflow(graph_def, layout='NCHW') when importing Tensorflow graph definition to nnvm. Does this mean that the nnvm will generate the graph according to the 'NCHW` layout?

  2. Then, when we call nnvm.compiler.build, it seems to me that the document wasn’t quite clear about what shape is exactly. If the original input to the graph is (1, 224, 224, 3) in NHWC layout, which one should we pass to the function:

  • shape={input: (1, 224, 224, 3)} or
  • shape={input: (1, 3, 224, 224)}?

Example code:

x = tf.placeholder(dtype=tf.float32, shape=(None, 224, 224, 3), name='input')
w = tf.constant(kernel, dtype=tf.float32, shape=(3, 3, 3, 1), name='w')
y = tf.nn.conv2d(input=x, filter=w, strides=[1, 1, 1, 1],
                 padding='SAME',
                 data_format='NHWC',
                 name='output')

# Assume we train the graph, and save the graph in '.pb' format.

# Load the graph from '.pb' file.
...

with tf.Graph().as_default() as graph:
    tf.import_graph_def(graph_def, name='')
    sym, params = nnvm.frontend.from_tensorflow(graph_def, layout='NCHW')

    shape_dict = {'input': (1, 3, 224, 224)}
    dtype_dict = {'input': 'float32'}

    graph, lib, params = nnvm.compiler.build(
        graph=sym,
        shape=shape_dict,
        dtype=dtype_dict,
        target=target,
        params=params,
        target_host=target_host)

    m = graph_runtime.create(graph, lib, ctx)
    m.set_input(**params)

#2

Sorry, Tensorflow frontend not support NCHW. If you pass NCHW, just means we will use NCHW for convolution, others will still be NHWC. So you will still pass (1, 224, 224, 3), not (1, 3, 224, 224). We have one RFC to discuss this, welcome any comment: https://github.com/dmlc/tvm/issues/2519


#3

Thanks. I am a very beginner but really into this right now and willing to contribute.

I am a little confused, please correct me if I’m wrong. So to clarify what you just explained.

  • First, if we pass NCHW to nnvm.frontend.from_tensorflow, it means the convolution will be calculated using NCHW format.
  • Second, the other parts of the graph will still be ‘NHWC’ as defined in the original Tensorflow graph.

#4

Yes. You could refer the RFC mentioned before to know more details.


#5

Yes, layout=NCHW is just to indicate using NCHW across NN operators. This allows compilation for CUDA target as CUDA doesn’t support NHWC now.