Conversion of PyTorch,pt model file into torchScript.ts file

username: Arpan_Gyawali
cooked: <p>I have a weight model named model.pt for brain segmentation from head CT scan.<br>
How can i convert this into torchscript file so that i can use the model for deployment,</p>
<p><strong>Network defn:</strong></p>
<pre><code class="lang-auto">3dUNet, 
in channel: 1(image), 
out channel: 2(brain label and background)
</code></pre>
<p><strong>Input defn:</strong></p>
<pre><code class="lang-auto"> "image": {
                "type": "image",
                "format": "hounsfield",
                "modality": "CT",
                "num_channels": 1,
                "spatial_shape": [
                    96,
                    96,
                    96
                ],
                "dtype": "float32",
                "value_range": [
                    0,
                    1
                ],
                "is_patch_data": true,
                "channel_def": {
                    "0": "image"
                }
            }
        },
</code></pre>
<p><strong>Train/val split:</strong> 13 images for training and 3 for validation</p>
<p><strong>Output defn:</strong></p>
<pre><code class="lang-auto">          "pred": {
              "type": "image",
              "format": "segmentation",
              "num_channels": 2,
              "spatial_shape": [
                  96,
                  96,
                  96
              ],
              "dtype": "float32",
              "value_range": [
                  0,
                  1
              ],
              "is_patch_data": true,
              "channel_def": {
                  "0": "background",
                  "1": "brain"
              }
          }
      
</code></pre>
<p>Now, how can i use tracing/scripting to convert into torchsctipt.<br>
Are these pieces of information enough?</p>
<p>I tried</p>
<pre><code class="lang-auto">import torch

model = torch.load('model/model.pt')

example = torch.rand(13, 96, 96, 96)

traced_script_module = torch.jit.script(model, (example))
torch.save(traced_script_module, "model/traced_resnet_model.ts")
</code></pre>
<p>I only used the model input size, I also tried <code>torch.jit.trace</code>. But both failed.<br>
Any help would be very appreciable.</p>

username: Paplham
cooked: <p>Could you provide more information on the <code>model.pt</code> file? Could you also provide us with the error? It is hard to tell what could be wrong if we do not see the error message.</p>
<p>Anyways, my guess would be that your <code>model.pt</code> file contains only the model weights, e.g., saved by</p>
<pre><code class="lang-auto">torch.save(model.state_dict(), PATH)
</code></pre>
<p>Is that assumption correct? If so, the file does not contain information about the actual model architecture, and only contains the trained weights. You should instantiate the model and then load the weights. Afterwards, you can attempt to convert it to jit script.</p>

username: Arpan_Gyawali
cooked: <p>Thanks <a class="mention" href="/u/paplham">@Paplham</a> for your reply<br>
Yes model.pt contains only the model weights and the error i got when i deployed the model using monai-deploy-app-sdk was:<br>
<code># ItemNotExistsError: A predictor of the model is not set.</code></p>
<p>If I instantiate the model and then load the weights. Then what should i do to convert it to jit script?</p>

username: Paplham
cooked: <p>First, you need to instantiate the model and load the weights. Please, see <a href="https://pytorch.org/tutorials/beginner/saving_loading_models.html" class="inline-onebox" rel="noopener nofollow ugc">Saving and Loading Models — PyTorch Tutorials 2.0.1+cu117 documentation</a> for this.</p>
<p>I will also provide you with an example below, where I will be using a <code>resnet18</code> from torchvision.</p>
<p>First, we instantiate the model:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torchvision
from torchvision.models import resnet18

model = resnet18(pretrained=False)
</code></pre>
<p>You can then train the model and afterwards, save the trained weights as:</p>
<pre><code class="lang-auto">torch.save(model.state_dict(), 'trained_weights.pt')
</code></pre>
<p>You have probably already done those steps. Now, we want to load the model and export it to torchscript. For this, you need to instantiate the model again and load the saved weights into it:</p>
<pre><code class="lang-auto">model = resnet18(pretrained=False)
model.load_state_dict(torch.load('trained_weights.pt'))
model.eval()
</code></pre>
<p>If your model uses standard PyTorch operations, exporting it should be as simple as calling:</p>
<pre><code class="lang-auto">model_scripted = torch.jit.script(model) # Export to TorchScript
torch.jit.save(model_scripted, 'exported_model.pt')
</code></pre>
<p>Please, see <a href="https://pytorch.org/docs/stable/generated/torch.jit.script.html#torch.jit.script" class="inline-onebox" rel="noopener nofollow ugc">torch.jit.script — PyTorch 2.0 documentation</a> for more details.</p>

username: Arpan_Gyawali
cooked: <aside class="quote no-group" data-username="Paplham" data-post="4" data-topic="185671">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/paplham/48/62872_2.png" class="avatar"> Paplham:</div>
<blockquote>
<p><code>torch.jit.save(model_scripted, 'exported_model.pt')</code></p>
</blockquote>
</aside>
<p>Thank you so much <a class="mention" href="/u/paplham">@Paplham</a><br>
I have done this</p>
<pre><code class="lang-auto">import torch
from monai.networks.nets import UNet

model = UNet(
            spatial_dims=3,
            in_channels=1,
            out_channels=2,
            channels=[16, 32, 64, 128, 256],
            strides=[2, 2, 2, 2],
            num_res_units=2,
            norm="batch",
        )

model.load_state_dict(torch.load('model/model.pt'))

model.eval()

# example = torch.rand(13, 96, 96, 96)

model_scripted = torch.jit.script(model)
torch.jit.save(model_scripted, 'model/torchscript_model.ts')
</code></pre>
<p>Dont we need the model input shape for passing into torch.script?</p>

username: Paplham
cooked: <p>Depends on the model that you are using. Generally, you do not need to provide the input shape. Perhaps, you are thinking of <a href="https://pytorch.org/docs/stable/generated/torch.jit.trace.html#torch.jit.trace" rel="noopener nofollow ugc">tracing</a>?</p>

username: Arpan_Gyawali
cooked: <p>Yes, what is the actual difference between scripting and tracing. Why does tracing need input shape and scripting donot?</p>

username: Paplham
cooked: <p>There is a long discussion to be had about that subject, if you are interested in it, a lot can be found in the <a href="https://pytorch.org/docs/stable/jit.html" rel="noopener nofollow ugc">documentation</a>.</p>
<p>However, basic intuition of the two is as follows:</p>
<p><strong>Scripting</strong> will inspect the <strong>source code</strong> and compile it, also saving the needed weights and possibly also functions of the model (if wrapped in an export decorator and possibly much more can be done!).</p>
<p><strong>Tracing</strong> is done by passing input to the model and <strong>remembering which operations</strong> were done to the input in order to get the output. These operations are then compiled and saved.</p>
<p>What is the practical difference? Well, if your model’s forward pass contains an <code>if</code> <code>else</code> statement, then what will happen if we script vs. trace?</p>
<p>If we trace the model, only one condition from the <code>if/else</code> will be executed and the exported model will always perform the same operation. If we script the model, the <code>if/else</code> condition will be checked and possibly different operations will be run.</p>
<p>There are more differences, but I think that the <code>if/else</code> example showcases well what the differences are.</p>

username: Arpan_Gyawali
cooked: <p>Thank you <a class="mention" href="/u/paplham">@Paplham</a> for your help. My issue is solved.</p>

username: Paplham
cooked: <p>Glad I could be of help <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>