Dynamically skip some front part of torchscript model

Jungmo_Ahn: <p>Hello</p>
<p>I’m using 1.5.0 version of PyTorch.<br>
I’d like to skip <em>i</em> numbers of layers in a Torchscript model.<br>
Below code is my trial. The method ‘skip’ takes input tensor x and skip point indicator <em>i</em>.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

class Model(torch.nn.Module):
        def __init__(self):
            super(Model, self).__init__()
            self.fc1 = nn.Linear(20,100)
            self.fc2 = nn.Linear(100, 100)
            self.fc3 = nn.Linear(100, 1)
            self.k = nn.ModuleList()
            self.k.append(self.fc1)
            self.k.append(self.fc2)
            self.k.append(self.fc3)

        @torch.jit.export
        def skip(self, x, i):
            for l in range(int(i), 3):
                x = self.k[l].forward(x)
            return x

        def forward(self, x):
            x = self.fc1(x)
            x = self.fc2(x)
            x = self.fc3(x)
            return x

model = Model().cuda()
x = torch.randn(20).cuda()
x2 = torch.randn(100).cuda() # Assume that x2 is the result of self.fc1 which is computed in advance.
idx = torch.tensor(1).cuda()
print(model.skip(x2,idx)) # This works

traced = torch.jit.trace(model, x) # Script
print(traced.skip(x2,idx)) # This does not works. Error occurs.
</code></pre>
<p>The error message is printed as below.</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "C:/Users/user/PycharmProjects/coin-modelgen/simple.py", line 33, in &lt;module&gt;
    traced = torch.jit.trace(model, x) # Script
  File "C:\Users\user\Anaconda3\envs\pytorch\lib\site-packages\torch\jit\__init__.py", line 875, in trace
    check_tolerance, _force_outplace, _module_class)
  File "C:\Users\user\Anaconda3\envs\pytorch\lib\site-packages\torch\jit\__init__.py", line 1021, in trace_module
    module = make_module(mod, _module_class, _compilation_unit)
  File "C:\Users\user\Anaconda3\envs\pytorch\lib\site-packages\torch\jit\__init__.py", line 716, in make_module
    return torch.jit._recursive.create_script_module(mod, make_stubs_from_exported_methods, share_types=False)
  File "C:\Users\user\Anaconda3\envs\pytorch\lib\site-packages\torch\jit\_recursive.py", line 305, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "C:\Users\user\Anaconda3\envs\pytorch\lib\site-packages\torch\jit\_recursive.py", line 361, in create_script_module_impl
    create_methods_from_stubs(concrete_type, stubs)
  File "C:\Users\user\Anaconda3\envs\pytorch\lib\site-packages\torch\jit\_recursive.py", line 279, in create_methods_from_stubs
    concrete_type._create_methods(defs, rcbs, defaults)
RuntimeError: 
Expected integer literal for index:
  File "C:/Users/user/PycharmProjects/coin-modelgen/simple.py", line 18
        def skip(self, x, i):
            for l in range(int(i), 3):
                x = self.k[l].forward(x)
                    ~~~~~~~~ &lt;--- HERE
            return x
</code></pre>
<p>Could someone can share any idea (or alternative approach) to solve this problem?</p>
<p>Moreover, I want to know that such dynamically skipping is also possible or not in c++ implementation using libtorch.</p>
<p>According to <a href="https://github.com/pytorch/pytorch/issues/47496" class="inline-onebox" rel="noopener nofollow ugc">torch.jit.script cannot index a ModuleList with an int returned from range() · Issue #47496 · pytorch/pytorch · GitHub</a>, I identified that the <em>i</em> must be selected in compile time not runtime. Thus…, is this meaning that there no way to dynamically skip the front part of the Torchscript model?</p>
<p>Or, is such naive below code a unique way?</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

class Model(torch.nn.Module):
        def __init__(self):
            super(Model, self).__init__()
            self.fc1 = nn.Linear(20,100)
            self.fc2 = nn.Linear(100, 100)
            self.fc3 = nn.Linear(100, 1)
            self.k = nn.ModuleList()
            self.k.append(self.fc1)
            self.k.append(self.fc2)
            self.k.append(self.fc3)
            # self.k = nn.Sequential(self.fc1,self.fc2,self.fc3 )

        def from1(self, x):
            for index, v in enumerate(self.k[1:]):
                x = v.forward(x)
            return x

        def from2(self, x):
            for index, v in enumerate(self.k[2:]):
                x = v.forward(x)
            return x

        @torch.jit.export
        def skip(self, x, i:int):
            if i == 1:
                x = self.from1(x)
            elif i == 2:
                x = self.from2(x)
            return x

        def forward(self, x):
            x = self.fc1(x)
            x = self.fc2(x)
            x = self.fc3(x)
            return x

model = Model().cuda()
m = torch.jit.script(Model()).cuda()

x = torch.randn(20).cuda()
x2 = torch.randn(100).cuda() # Assume that x is the result of self.fc1 which is computed in advance.
idx = 1
print(model.skip(x2,idx)) # This works
print(m.skip(x2,idx)) # This works
</code></pre>
<pre><code class="lang-auto">tensor([-0.3310], device='cuda:0', grad_fn=&lt;AddBackward0&gt;)
tensor([-0.2432], device='cuda:0', grad_fn=&lt;AddBackward0&gt;)
</code></pre>
<p>If this way is the unique way, why the two result are different?</p>
<p>Thanks a lot.</p>

SplitInfinity: <aside class="quote no-group" data-username="Jungmo_Ahn" data-post="1" data-topic="112442">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/41988e/48.png" class="avatar"> Jungmo_Ahn:</div>
<blockquote>
<p>Could someone can share any idea (or alternative approach) to solve this problem?</p>
</blockquote>
</aside>
<p>One thing you can do is iterate over the entire module list with <code>enumerate</code> and only call each module if <code>i</code> is greater than some limit <code>k</code>. Performance will be worse than if you skipped the first <code>k</code> layers completely.</p>
<aside class="quote no-group" data-username="Jungmo_Ahn" data-post="1" data-topic="112442">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/41988e/48.png" class="avatar"> Jungmo_Ahn:</div>
<blockquote>
<p>According to <a href="https://github.com/pytorch/pytorch/issues/47496" rel="noopener nofollow ugc">torch.jit.script cannot index a ModuleList with an int returned from range() · Issue #47496 · pytorch/pytorch · GitHub</a>, I identified that the <em>i</em> must be selected in compile time not runtime. Thus…, is this meaning that there no way to dynamically skip the front part of the Torchscript model?</p>
</blockquote>
</aside>
<p>You can dynamically skip the front part but only by annotating the LHS of your indexing statement with a module interface type. See the issue you linked for more details. I am currently working on extending this for <code>ModuleList</code>, it currently only works for <code>ModuleDict</code>.</p>
<aside class="quote no-group" data-username="Jungmo_Ahn" data-post="1" data-topic="112442">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/41988e/48.png" class="avatar"> Jungmo_Ahn:</div>
<blockquote>
<p>If this way is the unique way, why the two result are different?</p>
</blockquote>
</aside>
<p>I think it’s because <code>from1</code> and <code>from2</code> perform different computations?</p>