Failed to profile torchscript models

[terryzhao127]

I tried to apply this nice profiler to torchscript models but it failed:

Traceback (most recent call last):
  File "/home/guikarist/projects/test-torch/profiler/torchnnprofiler/test.py", line 24, in <module>
    with LayerProf(net) as prof:
  File "/home/guikarist/projects/test-torch/profiler/torchnnprofiler/nnprofiler/__init__.py", line 43, in __init__
    self._attach_backward_hook(name)
  File "/home/guikarist/projects/test-torch/profiler/torchnnprofiler/nnprofiler/__init__.py", line 84, in _attach_backward_hook
    self._register_hooks(name)
  File "/home/guikarist/projects/test-torch/profiler/torchnnprofiler/nnprofiler/__init__.py", line 77, in _register_hooks
    self._hook_handles.append(layer.register_forward_pre_hook(fw_pre_hook))
  File "/home/guikarist/.miniconda3/envs/pytorch-2.0.1/lib/python3.9/site-packages/torch/jit/_script.py", line 942, in fail
    raise RuntimeError(name + " is not supported on ScriptModules")
RuntimeError: register_forward_pre_hook is not supported on ScriptModules

Is there anyway to solve the problem? This issue may be related. And register_forward_pre_hook post-scripted models has not been supported yet: https://github.com/pytorch/pytorch/issues/29037#issuecomment-792868448

[kshitij12345]

@terryzhao127 thanks for raising the issue, but you are right that since ScriptModules don't support hooks, we can't leverage this implementation to profile those models 😓 (as of now).