4 – PyTorch Script Annotation V2

Well, sometimes your model might use forms of control flow that don’t actually work with this tracing method. So, for example, you might have some if statements in your forward method that depend on your input. So, things like this, they’re fairly common in natural language processing problems. So, there is a second way to convert your models to Tor script and this is through annotation. So, if you have a model like this. So, typically you would subclass it from torch.nn.module, create your parameters or whatever else you need, and then have some forward pass. Here, the forward pass has this If statement that depends on the input itself. So, here the tracing method that we used before isn’t going to work. Remember with that, what we did is we basically passed an example input through our network, and then in that way like traced out all the operations and built this graph, but with control flow like this then the graph that you’re going to build actually depends on the input that you put in. So, instead what we’re gonna do here is subclass from torch.jit.ScriptModule. So remember, when we used tracing that that actually returned a script module for us, but here we are creating our own script module. Now, with our forward method, we know with this control flow we can use a decorator, so torch.jit.ScriptMethod, and then it will appropriately convert this to a script module for us. So, the cool thing about this is that you previously defined your module like this, and you’ve got everything to work, you’ve trained it, and now you want to convert it for production. All you really need to do is subclass it from script module instead of module here, and then add this decorator. So, it’s basically just two changes to the code that you already have, and you’re ready to ship it to production after training. So now, with either method, tracing or annotation to get your script module, you can serialize it to a file which can then later be loaded into C++. So, to do this, you simply take your script module and use the save method, passing it in the file name or path where you want it to be saved, and then it will produce this file in your directory. With this file, you can load it into a C++ program, and use your model just as you would in Python.

%d 블로거가 이것을 좋아합니다: