YOLOX to TVM

[jlamperez]

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]

Hi @jlamperez ,

1. tvm discourage the use of aten::copy_, you can check this discussion: Adding aten::unsqueeze_ to PT Frontend apache/tvm#7231 (comment) for more details.
2. for aten::silu_, you can use something like below to resolve this ops:
   

   YOLOX/yolox/models/network_blocks.py

   Lines 9 to 14 in 902b372

   	class SiLU(nn.Module):
   	"""export-friendly version of nn.SiLU()"""
   	@staticmethod
   	def forward(x):
   	return x * torch.sigmoid(x)

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.