Nnvm CorrectLayout pass for NHWC


Hello! I’m trying to work with NNVM’s CorrectLayout pass for NHWC format. The computation graph is intialized to NCHW format. It looks like if I pass layout=‘NHWC’ to the build() function in the nnvm/python/nnvm/compiler/build_module.py, the CorrectLayout pass doesn’t change the node’s parameters, e.g. BatchNorm’s ‘axis’ parameter, which is set to 1 as default for NCHW, and therefore it doesn’t give the correct layouts for nodes in a graph.

How can I change a node’s parameters if necessary? Thanks in advance!


We have moved our Graph IR to Relay, so NNVM is deprecated.

In relay’s AlterOpLayout pass, we assume batch norm is expanded to several elemwise/broadcast ops, and then we can handle it correctly.
We do modify the attribute (a dangerous modification to const argument) for pooling operator, you can take it as an example if you want to modify the attribute of batch norm (https://github.com/dmlc/tvm/blob/51785062553b19a219b2639cd230b7da7455cd0d/src/relay/op/nn/pooling.cc#L46)

related discuss : [DISCUSS] and Documentation of AlterLayout Pass