Batchnorm op Fusion in TVM

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