search

zhanghang1989 / ResNeSt

ResNeSt: Split-Attention Networks
https://arxiv.org/abs/2004.08955
Apache License 2.0
3.24k stars 496 forks source link

Unable to export resnest model to torch script #64

Open lessw2020 opened 4 years ago

lessw2020 commented 4 years ago

Very impressed with the ResNest accuracy and want to move a model to production. However, in trying to script it, I hit the following issue:

`~\anaconda3\envs\fastai2\lib\site-packages\torch\jit_recursive.py in create_script_module_impl(nn_module, concrete_type, stubs_fn) 359 # Compile methods if necessary 360 if concrete_type not in concrete_type_store.methods_compiled: --> 361 create_methods_from_stubs(concrete_type, stubs) 362 torch._C._run_emit_module_hook(cpp_module) 363 concrete_type_store.methods_compiled.add(concrete_type)

~\anaconda3\envs\fastai2\lib\site-packages\torch\jit_recursive.py in create_methods_from_stubs(concrete_type, stubs) 277 rcbs = [m.resolution_callback for m in stubs] 278 defaults = [get_default_args(m.original_method) for m in stubs] --> 279 concrete_type._create_methods(defs, rcbs, defaults) 280 281 def create_script_module(nn_module, stubs_fn, share_types=True):

RuntimeError: Can't redefine method: forward on class: torch.resnest.torch.splat.SplAtConv2d

`

Is there any input on how to export to script so we can use in production?
Thanks very much!

zhanghang1989 commented 4 years ago

Sorry that I am not familiar with torch script.

lessw2020 commented 4 years ago

Thanks @zhanghang1989 - I'm going to run it in eager mode for now and may see if I can take a crack at the jit aspect next week. Congrats btw on the resnest model. I'm super impressed with the results from it!

lessw2020 commented 4 years ago

Hi @zhanghang1989 - good news in that @rwightman was able to make the fix to support JIT fix with ResNeST, and I've been able to export as JIT and run live on our server.

Here's the link to his version of your ResNest with JIT fix: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/resnest.py

I think the fix is adding a dummy self.bn2 in the split attention conv2d to keep JIT happy but you'll be better able to determine looking at the code or checking with Ross directly.

Anyway, thanks again for making ResNeST - it's accuracy is the best of the models I've tested in my work!