How to jit script a module that uses spectral norm?

user_10784842 0
Viewed 145

I am attempting to speed up my model training by compiling my model with the pytorch jit scripting feature. Unfortunately, my model needs to make use of Spectral Norm and it seems jit does not work with it:

In [2]: import torch
In [3]: c1 = torch.nn.utils.spectral_norm(torch.nn.Conv2d(1,2,3))
In [4]: c2 = torch.jit.script(c1)
Traceback (most recent call last):
  [...]
  File "[...]torch\jit\_recursive.py", line 761, in get_hook_stubs
    if pre_hook.__name__ in hook_map:
AttributeError: 'SpectralNorm' object has no attribute '__name__'

I noticed that there is now a new implementation of spectral norm, but it also does not compile successfully:

In [5]: c3 = torch.nn.utils.parametrizations.spectral_norm(torch.nn.Conv2d(1,2,3))
In [6]: c4 = torch.jit.script(c3)
Traceback (most recent call last):
  [...]
  File "[...]torch\jit\frontend.py", line 286, in __call__
    raise UnsupportedNodeError(ctx, node)
torch.jit.frontend.UnsupportedNodeError: GeneratorExp aren't supported:
  File "[...]torch\nn\utils\parametrize.py", line 262
            x = self[0](self.original)
        else:
            originals = (getattr(self, f"original{i}") for i in range(self.ntensors))
                        ~ <--- HERE
            x = self[0](*originals)
        # It's not possible to call self[1:] here, so we have to be a bit more cryptic

I am using pytorch 1.11.0+cu113 from pip, on Windows.

Is there a workaround for this currently? Or am I doing something wrong?

pytorch
0 Answers
Related