You are viewing a plain text version of this content. The canonical link for it is here.
Posted to discuss-archive@tvm.apache.org by Aakanksha via Apache TVM Discuss <no...@discuss.tvm.ai> on 2022/03/23 07:52:48 UTC
[Apache TVM Discuss] [Questions] Batchnorm op Fusion in TVM
Dear All,
I am looking for a set transformation passes in TVM that helps in fusing/folding the Batchnorm ops into the previous or the next convolution-like layers.
**My expectation :**
* **before batchnorm fold** : conv2d -> bias_add -> batch_norm
* **after batchnorm fold** : conv2d *(possibly changed weights)* -> bias_add *(possibly changed bias)*
The mathematics for this above transformation can be understood from below set of images.
(source of the below snapshots - https://nenadmarkus.com/p/fusing-batchnorm-and-conv/)
As of now I have been able to find SimplifyInference pass in TVM that is related to simplification of Batchnorm op.
But, from what I understand about this pass is that it separates out the constant terms from the batchnorm operation and folds it but the ops corresponding to the terms with data involvement are still present in the relay graph as two basic operations "Multiply" and "Add".
**Add ( Multiply ( data, scale ) , shift )**
where :
scale is [gamma/sqrt(running_variance + epsilon)]
shift is [{-(running_mean * gamma)/(sqrt(running_variance + epsilon)}+beta]
I have applied "FoldScaleAxis" and "FoldConstant" pass in the same order after the Simplify inference pass which are not helpful in what I expect to achieve.
Can someone suggest if TVM has any other set of transformation passes working at the relay level that can help me get the expected Batchnorm fuse/fold transformation over the relay graph?
Thanks!!
cc: @mbrookhart @mbs-octoml
---
[Visit Topic](https://discuss.tvm.apache.org/t/batchnorm-op-fusion-in-tvm/12391/1) to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/d7db6843094e434a76d589a2ad9b8181569ddad5502763877a8fb846aa6ad71f).
[Apache TVM Discuss] [Questions] Batchnorm op Fusion in TVM
Posted by Aakanksha via Apache TVM Discuss <no...@discuss.tvm.ai>.
Hi @masahi ! Thanks for the quick response. I tried the sequence of passes you suggested but still seeing the same effect, i.e., multiply and add ops in place of Batchnorm op.
cc: @mbrookhart
---
[Visit Topic](https://discuss.tvm.apache.org/t/batchnorm-op-fusion-in-tvm/12391/3) to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/a194d2876e3c82f931d8ff1fe0fc6bc219645b161f65de71f72639391553bd2d).
[Apache TVM Discuss] [Questions] Batchnorm op Fusion in TVM
Posted by masahi via Apache TVM Discuss <no...@discuss.tvm.ai>.
You need to apply `FoldScaleAxis` after `FoldConstant`. See https://github.com/apache/tvm/blob/ac6607282e080dc15cce7d9cf565f5d390ba0f16/tests/python/relay/test_pass_fold_constant.py#L316-L323
---
[Visit Topic](https://discuss.tvm.apache.org/t/batchnorm-op-fusion-in-tvm/12391/2) to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/50819738954530bc1a4eab0a8f5d3958d750da777a546e16936b8abb5b872ec7).
[Apache TVM Discuss] [Questions] Batchnorm op Fusion in TVM
Posted by masahi via Apache TVM Discuss <no...@discuss.tvm.ai>.
Looks like `SimplifyExpr` doesn't support folding `bias_add` and `bias`, see https://github.com/apache/tvm/blob/6942b3660df3551a3a9a86c2faba834d366a2a7e/src/relay/transforms/simplify_expr.cc#L651-L652. So both of cases don't work unless you modify that pass. But I recommend not depending on `bias_add` as explained below.
[quote="aakaverm-quic, post:9, topic:12391"]
As I mentioned in my original question I would need to preserve the conv2d and bias_add ops after batchnorm fold. It is more of a pattern matching requirement rather than an optimization one.
[/quote]
I highly suggest modifying your pattern to support both `bias_add` and `add` as in https://github.com/apache/tvm/blob/7fd73b2663ae33d341ab09834f215285eb9bd136/python/tvm/relay/op/contrib/cutlass.py#L45
Frontends are not consistent in which ops to generate for bias addition.
And I recommend experimenting with simple test cases like https://github.com/apache/tvm/blob/ac6607282e080dc15cce7d9cf565f5d390ba0f16/tests/python/relay/test_pass_fold_constant.py#L305 rather than the whole resnet50 from the beginning.
---
[Visit Topic](https://discuss.tvm.apache.org/t/batchnorm-op-fusion-in-tvm/12391/11) to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/7962288d6c528b68e4720ba60b7c48b30b464c26528e4b3e8b1f413306aba6e7).
[Apache TVM Discuss] [Questions] Batchnorm op Fusion in TVM
Posted by Aakanksha via Apache TVM Discuss <no...@discuss.tvm.ai>.
Hi @masahi , I am not quite clear regarding this bias_add and add op folding that you mentioned.
So, what I intend to achieve and what I assume you are also implying above is as follows:
case 1:
>**before:** conv2d -> bias_add -> add (shift from batchnorm) is transformed to:
>
> **after transform:** conv2d -> bias_add ( bias values are changed by add op folding into bias_add op)
case 2:
> **before:** conv2d -> add (shift from batchnorm) is transformed to:
>
>**after transform:** conv2d -> biad_add (add op expressed as bias _add op)
Can you please confirm once if this is what you also meant above.
Thanks for the catch.. yeah I tried now with SimplifyExpr but not seeing the desired changes in the relay graph.
I shall see if I can find some helpful TVM transformation passes sequence in the cases 1 & 2.
---
[Visit Topic](https://discuss.tvm.apache.org/t/batchnorm-op-fusion-in-tvm/12391/9) to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/de1e4a219480588be301ae94a9bf240b8aaa85d7b3b4971f20198dd5d2512070).
[Apache TVM Discuss] [Questions] Batchnorm op Fusion in TVM
Posted by Ligeng Zhu via Apache TVM Discuss <no...@discuss.tvm.ai>.
Either case should be possible with tvm transformation pass.
If you need to preserve `conv2d` and `biad_add`, an alternative way to achieve this is to perform such transform in original models before exporting to TVM.
For example, I used to use this [script](https://gist.github.com/Lyken17/deb98385a06ae67ce1252c1a17ad181d) to fuse bn into conv.
---
[Visit Topic](https://discuss.tvm.apache.org/t/batchnorm-op-fusion-in-tvm/12391/10) to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/30f2f6ae7eed4ad3d2ddc030226a7b3676caab0e43dbaf1830a30b5311845634).
[Apache TVM Discuss] [Questions] Batchnorm op Fusion in TVM
Posted by Aakanksha via Apache TVM Discuss <no...@discuss.tvm.ai>.
I hadn't run `bind_param_by_name`. I tried it now. I am now not seeing multiply ops however I still see add ops in place of batchnorm ops. The script I am using is given below.
Thanks @masahi !
> import onnx
>
> import tvm
>
> from tvm import relay
>
>from tvm.relay.build_module import bind_params_by_name
>
>import os
>
> os.system("wget https://s3.amazonaws.com/onnx-model-zoo/resnet/resnet50v1/resnet50v1.onnx")
>
> dtype_dict = {"data": "float32"}
>
> shape_dict = {"data": [1,3,224,224]}
>
> onnx_model = onnx.load('resnet50v1.onnx')
>
> mod, params = relay.frontend.from_onnx(onnx_model, shape_dict, freeze_params=True)
>
> print(mod)
>
> mod["main"] = bind_params_by_name(mod["main"], params)
>
>with tvm.transform.PassContext(opt_level=3):
>
> seq1 = tvm.transform.Sequential(
>
> [relay.transform.InferType(),
> relay.transform.SimplifyInference(),
> relay.transform.FoldConstant(),
> relay.transform.FoldScaleAxis(),
> ])
>
>
> mod = seq1(mod)
>
> print(mod)
cc: @mbrookhart
---
[Visit Topic](https://discuss.tvm.apache.org/t/batchnorm-op-fusion-in-tvm/12391/5) to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/f31777cf782380600375729307244cda71f52a90d205b72f63cd39a9821ff3af).
[Apache TVM Discuss] [Questions] Batchnorm op Fusion in TVM
Posted by masahi via Apache TVM Discuss <no...@discuss.tvm.ai>.
Having `add` there is expected since batch norm has a shift by constant. But the idea is that the new `add` can be folded into conv2d bias add. `SimplifyExpr` pass finds such two consecutive `add` with constant rhs, and fold them into one `add`.
---
[Visit Topic](https://discuss.tvm.apache.org/t/batchnorm-op-fusion-in-tvm/12391/6) to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/f27c8b484cdb0e889867360eb46e853ae3137d364dfc07e2541316bfb22201d7).
[Apache TVM Discuss] [Questions] Batchnorm op Fusion in TVM
Posted by masahi via Apache TVM Discuss <no...@discuss.tvm.ai>.
Note that you are using `SimplifyInfernce` twice, but you want to replace the second one with `SimplifyExpr`.
But right, it seems `bias_add` and `add` are not folded. It seems `relay.transform.CanonicalizeOps()` converts `bias_add` to `add`, so you want to call it before `SimplifyExpr`. I tried your script but I didn't get satisfying output after brief attempt. Maybe you need to play around it a bit more. If you believe there is a bug / missing functionality, welcome to open a github issue.
---
[Visit Topic](https://discuss.tvm.apache.org/t/batchnorm-op-fusion-in-tvm/12391/8) to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/755acbdd7f0b8091a486e0ebf6992d52f12f73840f7401f9bdb49e5a8f973f36).
[Apache TVM Discuss] [Questions] Batchnorm op Fusion in TVM
Posted by masahi via Apache TVM Discuss <no...@discuss.tvm.ai>.
Have you run `bind_param_by_name`? https://github.com/apache/tvm/blob/ac6607282e080dc15cce7d9cf565f5d390ba0f16/tests/python/relay/test_pass_fold_constant.py#L341
---
[Visit Topic](https://discuss.tvm.apache.org/t/batchnorm-op-fusion-in-tvm/12391/4) to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/d14f081b54b4148cfaf893ed286a4a9d0b5dde540348667f27996dd34871268c).
[Apache TVM Discuss] [Questions] Batchnorm op Fusion in TVM
Posted by Aakanksha via Apache TVM Discuss <no...@discuss.tvm.ai>.
Thanks @masahi.
Okay so, I tried with this sequence of passes:
>seq1 = tvm.transform.Sequential(
>
> [relay.transform.InferType(),
> relay.transform.SimplifyInference(),
> relay.transform.FoldConstant(),
> relay.transform.FoldScaleAxis(),
> relay.transform.SimplifyInference(),
> relay.transform.FoldConstant()
> ])
I get "add" ops as it is; they are not getting folded to the preceding conv2d's bias.
Also, if suppose there is no bias_add corresponding to a conv2d but batchnorm is there so after folding the batchnorm will a new bias_add op be created eventually to adjust the shift or the shift will remain as an add op in that case?
---
[Visit Topic](https://discuss.tvm.apache.org/t/batchnorm-op-fusion-in-tvm/12391/7) to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/6fbd906289a1e5b6a84436c58756c459409e123ddf0ed3f1f09c650e060223c1).