What is the meaning of @weak_module?

I saw in the code @weak_module decorator. Any explanation on what this means actually?

This code can help link

I saw this code, still the comments are missing for @weak_module.
My current understanding is that @weak_module should run faster than non weak one.

On line 99 there is a function for weak_module

From my understandingof the code, this is needed for torch.jit only to create a static computation graph, meaning that the function calls to/from this class will be lazily compiled and inlined

1 Like

Thanks @Kushaj great tips for weak_module.

Thanks @justusschock, nice tip, I know Python is interpreted language, just like Java, C# and PHP, and everything is JIT (just in time).

Why do we specifically need torch.jit? For introspection? Debugging?

The “Problem” here is, that PyTorch creates a dynamic computation graph at runtime, which is great for prototyping and allowing us to use python’s native flow-control operations, but unfortunately it makes it very hard to deploy models.

To tackle this problem, torch.jit traces the dynamically created graph and compiles it into a static version, which can be deployed e.g. via C++ without the need to reimplement the model’s logic.


Perfect feedback @justusschock, I see what you mean. I just found this:

TorchScript is a way to create serializable and optimizable models from PyTorch code. Any code written in TorchScript can be saved from your Python process and loaded in a process where there is no Python dependency.

We provide tools to incrementally transition a model from being a pure Python program to a TorchScript program that can be run independently from Python, for instance, in a standalone C++ program. This makes it possible to train models in PyTorch using familiar tools and then export the model to a production environment where it is not a good idea to run models as Python programs for performance and multi-threading reasons.

(from the official PyTorch documentation)

** Now I wonder:

We know that PyTorch on CUDA actually runs on C++. If this is true, then we don’t need to speed up.

One could get the idea if you run PyTorch code (Python) then you are slow. I bet this is not the case. Any clarification would be excellent.

You are not slow. You get the full performance of the optimized CUDA and C++ Code. The main problem lies within your graph/network definition:

You can achieve the same goal in different ways, each of them having it’s own (more or less efficient) graph definition. Theoretically it is possible to gain some speed by fusing kernels, optimizing the graph etc. But I must admit, that I don’t know whether there are any of these optimizations already implemented.

In general, running python code gives you a slight overhead (even if the backend is CUDA), because the python interpreter has some stuff like the GIL and some other book-keeping stuff, which is done at compile/linking time for pure C++ programs.

1 Like

I think I understood what you meant the very last time. You mentioned the deployment.
Actually at deployment your model is laying on some web server (possible Linux) and would be ideal to run on C++ to gain some speed over Python or PyPy.

Still, these speed improvements are not so crucial for most of the simple problems, but I can imagine there are some moments when speed actually matter.

I’d say, if you stick with a single thread, you would be fine with using python in most cases, but since you probably want to use multithreading (to have some background workers in other threads or simply forward the model in multiple threads), you would have to switch to multiprocessing in python, which has some hard limitations due to the GIL.

Also it is not that easy to setup python on most devices (like mobile or embedded devices, but it is virtually always possible to compile binaries [if you put enough effort to it]).

1 Like

Thanks @justusschock for the details. Your answer helped me to understand the way how to think on @weak_module decorator. Even though I don’t understand JIT I accepted your answer.