linghu8812 / tensorrt_inference

697 stars 205 forks source link

[ScaledYOLOv4] ONNX export failure on Google Colab runtimes #44

Open blah12629 opened 3 years ago

blah12629 commented 3 years ago

Here is the Colab link: https://colab.research.google.com/drive/1-Xn99ZMBHTj1h5Yo8Xjz9TiWbFwLEP41

I have followed the instructions, but it gives me the error ModuleNotFoundError: No module named 'mc'

I installed mish_cuda as mc and renamed ./build/lib.linux-x86_64-3.6/mish_cuda/to ./build/lib/mish_mish but then I get the error ONNX export failure: ONNX export failed: Couldn't export Python operator MishCudaFunction

I tried the workaround on here https://github.com/WongKinYiu/ScaledYOLOv4/issues/39#issuecomment-732036944 by changing from mc.build.lib.mish_mish import MishCuda as Mish to

 class Mish(nn.Module):
   def __init__(self):
       super().__init__()
   def forward(self,x):
       x = x * (torch.tanh(F.softplus(x)))
       return x

but it still gives ONNX export failure: ONNX export failed: Couldn't export Python operator MishCudaFunction.

I also tried replacing from mc.build.lib.mish_mish import MishCuda as Mish to

import importlib.util
spec = importlib.util.spec_from_file_location( \
  'mish_cuda', '/usr/local/lib/python3.6/dist-packages/mish_cuda-0.0.3-py3.6-linux-x86_64.egg/mish_cuda/__init__.py')
mish_cuda = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mish_cuda)
Mish = mish_cuda.MishCuda

but it still gave the same error.

bobbilichandu commented 3 years ago

@blah12629 were you able to figure this out?