How to script the model and lazily load parameters at the same time?

username: siahuat0727
cooked: <p>Hi, I am working on a tool to make model only load parameters when needed to reduce peak memory (<a href="https://github.com/siahuat0727/pylomin" rel="noopener nofollow ugc">link</a>). And I would like to take advantage of the <code>torch.jit</code> optimization but it failed. Here is a reproducible code:</p>
<pre><code class="lang-python">import functools
import torch
import torch.nn as nn
import torch.nn.functional as F  # for torch.jit.script


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(32, 32)
        self.fc2 = nn.Linear(32, 32)

    def forward(self, input):
        out = self.fc1(input)
        out = self.fc2(out)
        return out


def release_weights(module, _input=None, _output=None):
    module.to('cpu')


def load_weights(module, _input=None):
    module.to('cuda')


def lazy_loading(func, module):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        load_weights(module)
        res = func(*args, **kwargs)
        release_weights(module)
        return res
    return wrapper

def main():
    model = Net().eval()
    data = torch.rand(1, 32)

    for module in [model.fc1, model.fc2]:
        # method 1: decorator
        module.forward = lazy_loading(module.forward, module)

        # method 2: hook
        # module.register_forward_hook(release_weights)
        # module.register_forward_pre_hook(load_weights)

    data = data.cuda()

    with torch.no_grad():
        model = torch.jit.script(model)
        output = model(data)


main()
</code></pre>
<p>I’ve tried:</p>
<ol>
<li>For adding decorators (before/after <code>torch.jit.script</code>), the decorator seems to be ignored.</li>
<li>For adding hooks before <code>torch.jit.script</code>, it comes  <a href="https://discuss.pytorch.org/t/type-mismatch-with-hooks-input-when-using-jit/137783">Type mismatch with hooks</a>.</li>
<li>For adding hooks after <code>torch.jit.script</code>, it raises <code>RuntimeError: register_forward_hook is not supported on ScriptModules</code>.</li>
</ol>
<p>I wonder is it possible to script the model and lazily load parameters at the same time? (without directly modifying model internal code)<br>
Any feedback would be greatly appreciated. Thanks!</p>
<p>torch version: 1.10.1</p>

username: siahuat0727
cooked: <p>With method 2 + type hints (<a href="https://discuss.pytorch.org/t/type-mismatch-with-hooks-input-when-using-jit/137783/4">link</a>),</p>
<pre><code class="lang-python">...
def release_weights(module, _input: Tuple[torch.Tensor], _output):
    module.weight.data = module.weight.to('cpu')
...
</code></pre>
<p>it comes:</p>
<pre><code class="lang-nohighlight">RuntimeError:
Tried to set an attribute: data on a non-class: Tensor:
  File "/home/siahuat0727/test.py", line 22
def release_weights(module, _input: Tuple[torch.Tensor], _output):
    module.weight.data = module.weight.to('cpu')
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;--- HERE
</code></pre>
<p>So the main issue would be: <strong>Is it possible to move ScriptModule parameters across devices?</strong></p>
<p>Here’s a similar issue <a href="https://github.com/pytorch/pytorch/issues/28267" rel="noopener nofollow ugc">Problem with jit TorchScript while copying data between GRUs · Issue #28267 · pytorch/pytorch (github.com)</a>.</p>