Megvii-BaseDetection / YOLOX

YOLOX is a high-performance anchor-free YOLO, exceeding yolov3~v5 with MegEngine, ONNX, TensorRT, ncnn, and OpenVINO supported. Documentation: https://yolox.readthedocs.io/
Apache License 2.0
9.16k stars 2.16k forks source link

YOLOX to TVM #411

Open jlamperez opened 2 years ago

jlamperez commented 2 years ago

Hi,

Thank you for this nice repo.

I am trying to export pytorch model to TVM with this snippet of code

import torch
import tvm.relay as relay

from yolox.exp import get_exp

exp = get_exp(None, "yolox-nano")
exp.test_conf = 0.25
exp.nmsthre = 0.45
exp.test_size = (416, 416)
# exp

torch_model = exp.get_model()

ckpt_file = f"/tmp/YOLOX/{yolox_model.value}.pth"
ckpt = torch.load(ckpt_file, map_location="cpu")
torch_model.load_state_dict(ckpt["model"])
torch_model = torch_model.eval()

input_shape = [1, 3, 416, 416]
input_data = torch.randn(shape_dict["inputs"])
scripted_model = torch.jit.trace(torch_model, input_data).eval()

mod, params = relay.frontend.from_pytorch(scripted_model, shape_dict)

but I am having the next problem

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-54-fa43d6ab420b> in <module>
     22 scripted_model = torch.jit.trace(torch_model, input_data).eval()
     23 
---> 24 mod, params = relay.frontend.from_pytorch(scripted_model, shape_dict)
     25 
     26 # mod = relay.transform.InferType()(mod)

/opt/vitis_ai/conda/envs/vitis-ai-tensorflow/lib/python3.6/site-packages/tvm-0.8.dev1186+g40d5193a9-py3.6-linux-x86_64.egg/tvm/relay/frontend/pytorch.py in from_pytorch(script_module, input_infos, custom_convert_map, default_dtype)
   3289 
   3290     op_names = get_all_op_names(graph)
-> 3291     converter.report_missing_conversion(op_names)
   3292 
   3293     is_module = isinstance(script_module, torch.jit.ScriptModule)

/opt/vitis_ai/conda/envs/vitis-ai-tensorflow/lib/python3.6/site-packages/tvm-0.8.dev1186+g40d5193a9-py3.6-linux-x86_64.egg/tvm/relay/frontend/pytorch.py in report_missing_conversion(self, op_names)
   2549         if missing:
   2550             msg = "The following operators are not implemented: {}".format(missing)
-> 2551             raise NotImplementedError(msg)
   2552 
   2553     def convert_block(self, block, outputs):

NotImplementedError: The following operators are not implemented: ['aten::silu_', 'aten::copy_']

How these operators (silu and copy) can be implemented in TVM?

Best regards!

zhiqwang commented 2 years ago

Hi @jlamperez ,

  1. tvm discourage the use of aten::copy_, you can check this discussion: https://github.com/apache/tvm/pull/7231#issuecomment-756693366 for more details.
  2. for aten::silu_, you can use something like below to resolve this ops: https://github.com/Megvii-BaseDetection/YOLOX/blob/902b372be820a738ee95c55da2a957c5c2161536/yolox/models/network_blocks.py#L9-L14

BTW, I've done some experiments on YOLOv5 with tvm here, and the differences between this two models are relatively insignificant, so I believe with some minor changes we can also use tvm on YOLOX.