[FR] jit.trace module methods / module as an input

[ssnl]

Context: #17583

We killed support for tracing functions that reference tensor(s) requiring grad. This effectively makes it impossible to trace a function that computes on model parameters (e.g., compute the spectral norms of all conv weights in a CNN), unless we represent the function as another module. Yet using another module isn't always ideal, because then the two modules would need to share parameters (uhh).

An alternative workaround is to just put every parameter as the traced function's input, and then use another wrapper on top of the returned callable. This results in rather ugly and unreadable code.

So here is my FR: supporting either

• Tracing an nn.Module method, or
• Tracing with an nn.Module as an input

Really these two can maybe be viewed as the same thing because a method is just a bounded function with an input being the nn.Module object, i.e., self.

Here is the proposed API:

class Network(nn.Module):
  def __init__(self):
    self.conv = nn.Conv(3, 3, 3, 1)

  def forward(self, x):
    return self.conv(x).relu_()

  @jit.trace
  def estimate_spectral_norm(self, x):
    w = self.conv.weight
    # power iteration or really solve it or whatever
    ...

[suo]

@zdevito this would not be so hard if modules were 1st class, right?

[zdevito]

@ssnl I am not sure I understand the issue right, can you post an example failure. I want to make sure there isn't something we can do now.

[ssnl]

@zdevito

import torch
import torch.nn as nn


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 3, 3)

        # Want to trace self.weighted_kernel_sum

        # Method 1:
        #   torch.jit.trace(self.weighted_kernel_sum, (torch.randn(3, 3, 3, 3),))
        # Error:    
        #   RuntimeError: Cannot insert a Tensor that requires grad as a constant. 
        #                 Consider making it a parameter or input, or detaching the gradient

        # Method 2:
        #
        #   torch.jit.trace(Net.weighted_kernel_sum, (self, torch.randn(3, 3, 3, 3),))
        # Error:
        #   RuntimeError: Only tensors and (possibly nested) tuples of tensors 
        #   are supported as inputs or outputs of traced functions, but instead 
        #   got value of type Net.
        #   Value: Net(
        #     (conv): Conv2d(3, 3, kernel_size=(3, 3), stride=(1, 1))
        #   )

    def forward(self, x):
        return self.conv(x)

    def weighted_kernel_sum(self, weight):  # I want to trace this thing
        return (weight * self.conv.weight).sum()


Net()

[ssnl]

Method 1 used to work, until some time around January.

[ssnl]

@zdevito Does the above example make sense to you? :)

[zdevito]

Oh, sorry. I missed the reply a week ago :( Yes, the example makes sense. I think I understand the problem now. We can make this work.

[zdevito]

Here is my proposed fix: #19070

[suo]

Closing this as it's been folded into #19070. Feel free to reopen if you think the proposed fix is not enough