Pytorch JIT availability

Hi, I have a rather complex network structure that has a lot of python code in the forward pass. When I run it, the GPU is utilized only 50% while the CPU utilization is 600%. My guess is that the python code is the bottleneck and the whole thing is CPU bound, so I can’t use multiple GPUs available in my system to speed it up.

I was wondering if I could speed it up somehow? I am waiting for the JIT availability and would want to try it even in pre-release state. Are there any instructions on how to compile/use pytorch with JIT?

Thank you!


if you are on v0.2.0, try the master branch of pytorch, we’ve made most of the CPU bound models speedup quite a bit.

1 Like

After compiling the latest code, it runs much faster, consuming less CPU. Still when I run on a system with a faster performing single thread intel CPU, while using the exact same software and GPU hardware (1080ti) , I get a significant performance gain. This, plus the fact that nvidia-smi shows less than 50% GPU utilization makes me think that performance could still be improved. Is the JIT being employed transparently or perhaps I need to explicitly tell it to use JIT?

Thank you for the great tool, I love pytorch with its superior debugging and great flexibility!