Pytorch:How to change weights layout from NCHW to NHWC

Questions
#1

My device need the weights and input layout NHWC,but the pytorch model layout is NCHW.

I want to change weights layout from NCHW to NHWC , and I came up with two ways:

  1. In the TVM Relay,add transform layout before convolution.But this operation is too time consuming, and every time you run the network, you need to transform it again.

  2. A better way is,the parameter layout is converted to NHWC format in advance.But when I artificially transformed the layout in Pytorch model, here were some mistakes:

    • transform layout before jit.trance()

      model = weights_layout_NCHW2NHWnC(model)
model= torch.jit.trace(model, input_data).eval()

    The error is :

     Given groups=1, weight of size [64, 7, 7, 3], expected input[1, 224, 224, 3] to have 7 channels, but got 224 channels instead

    • transform layout after jit.trance() before relay.frontend.from_pytorch()

      model= torch.jit.trace(model, input_data).eval()
model = weights_layout_NCHW2NHWnC(model)
mod, params = relay.frontend.from_pytorch(model, shape_list)

    The error is :

    Shapes of input list and information in the graph do not match

This means that I can’t modify Pytorch model weights layout directly, Otherwise, there will be an error in the conversion TVM Relay.

Can you give me some suggestions to achieve my goal?

Thank you very much.

1 Like
#2

My device is embedded low power devices,so I don’t want these redundant operations to run on my device. I want to preprocess them on the host side as much as possible.

#3

You can use ConvertLayout pass for that purpose. This is how I transform NCHW pytorch models to NHWC, for example. torchscript-to-tvm/detr_test.py at master · masahi/torchscript-to-tvm · GitHub

After that, when you call relay.build(...) on the transformed module, parameters will also be transformed to NHWC at compile time. You can check the output of relay.build(...) to see that returned parameters are in NHWC (HWIO to be precise).

#4

[Relay] A set of utilities that allows a model to be run efficiently on tensorcores. by jwfromm · Pull Request #6748 · apache/tvm · GitHub This PR maybe help you!

#5

Thank for your reply.

I have a question that when I use this code :

desired_layouts = {'nn.conv2d': ['NHWC', 'default']}
# Convert the layout to NCHW
# RemoveUnunsedFunctions is used to clean up the graph.
seq = tvm.transform.Sequential([relay.transform.RemoveUnusedFunctions(),
                                relay.transform.ConvertLayout(desired_layouts)])
with tvm.transform.PassContext(opt_level=3):
    mod = seq(mod)
print(mod)

I got the following output (part):

%1 = layout_transform(%conv1.0.weight, src_layout="OIHW", dst_layout="HWIO") /* ty=Tensor[(7, 7, 3, 64), int8] */;
%2 = nn.conv2d(%0, %1, strides=[2, 2], padding=[3, 3, 3, 3], channels=64, kernel_size=[7, 7], data_layout="NHWC", kernel_layout="HWIO") /* ty=Tensor[(1, 112, 112, 64), int8] */;

Is the layout_transform() in the first row to be repeated every time the network is run? Want I mean is it necessary to convert layout a network weights every time you run a network?

It should only be done once.

I found after call relay.build(...) ,the return output parameters shape has been changed .So what is the use of layout_transform() in Relay graph?

#6

You don’t have to worry about layout transform on weight. So ConvertLayout pass inserts layout_transform to input and weights. Since weight is constant, we can do layout_transform(weight) at compile time, and that’s why you see weight shape has been changed after relay.build.

You can use opt_mod, opt_params = relay.optimize(mod, params, target) to see what the optimized graph looks like. There should be no layout_transform(weight). This is what relay.build() uses.

#7

Thank you for answering my questions.It’s very useful and answers my doubts.

I have another question .I have my own data layout,the weights layout is not the NHWC but the (N/16)HW(16n)C.

How can I customize the data layout?

Is there a document link to learn how to customize the data layout?

Thanks.

#8

I guess there is currently no inference strategy suitable for NHW16nC in TVM, which will make the inference process difficult if you have to keep your data layout :thinking:.

#9

My device only supports this layout, so I must have to change it. :rofl: