[AutoTVM] Tuning fails for an NHWC network on Arm CPU

[mbaret]

I'm trying to tune NHWC networks produced by the tflite frontend, but they fail with the error AssertionError: only support NCHW/HWCN currently. This suggests to me the AlterOpLayout pass is not happening during task extraction. I can force the AlterOpLayout pass to happen using the following patch:

--- a/python/tvm/autotvm/task/relay_integration.py
+++ b/python/tvm/autotvm/task/relay_integration.py
@@ -49,7 +49,9 @@ def _lower(func,
                 grc = graph_runtime_codegen.GraphRuntimeCodegen(None, target)
                 return grc.codegen(mod["main"])
     # default case
-    mod, _ = relay.optimize(func, target, params)
+    with relay.build_config(opt_level=3):
+        mod, _ = relay.optimize(func, target, params)
+
     grc = graph_runtime_codegen.GraphRuntimeCodegen(None, target)
     return grc.codegen(mod["main"])

which does fix this issue. Am I doing something wrong in AutoTVM or is this patch necessary?

[jwfromm]

You're not doing anything wrong. Not supporting NHWC layouts on Arm CPUs is a known issue and something we should have patched in soon.

[mbaret]

I think even if NHWC isn't supported, the auto-tuning shouldn't fail as a result. We can still compile networks that are initially in NHWC format as AlterOpLayout converts them to NCHW. I think AutoTVM should also do this instead of just failing.

[comaniac]

AutoTVM does not compile the model but just tune the ops extracted from the model. Accordingly, failing when the ops are not supported for tuning is reasonable.

Also please open such discussion topics in discuss.

[mbaret]

It doesn't compile, but it does run optimize which should invoke the appropriate passes to transform the graph. Is it a design choice for AutoTVM not to apply this pass? If so I'm happy to open a discuss post, as this is quite unintuitive behaviour.

[comaniac]

Not sure if that workaround works for all cases, but you are welcome to open an RFC to collect opinions from others.

[FrozenGene]

Please keep an eye of PR #3859 , which could solve TFLite NHWC auto tune on ARM CPU problem.

[anijain2305]

I agree that AlterOpLayout and AutoTVM do not have the most intuitive design/user experience yet.

For now, you can choose either of the 3 options

1. Use the [TOPI][AutoTVM] NHWC conv2d templates for ARM #3859 to use NHWC schedule. However, it might not give you the best performance because for now NCHW perf is somewhat better than NHWC.
2. Call AlterOpLayout pass the way you did before extracting tasks.
3. In general, use the [Relay] ConvertLayout pass. #4335 ConvertLayout to convert the whole graph to NCHW and then call AutoTVM extraction. Similarly, after AutoTVM is done, before calling the relay.build, call ConvertLayout.

[tqchen]

close for now due to inactive status, and there are related #3859 feel free to open new thread on the discuss forum