✨[Feature] support other methods compile besides 'forward' method

[lippman1125]

TRTorch now can only compile and keep forward method of torchscript. It is not transparency to users. If users exports other methods in torchscript, and their program is based on these methods。

[narendasan]

Can you provide a code example on how you would want this to work?

[lippman1125]

Can you provide a code example on how you would want this to work?

An example as follow, I want keep method 'get_length', but after compiled by using TRTorch, only 'forward' is reserved. If I want to call 'get_length' in my program, it will crash.

class LeNetFeatExtractor(nn.Module):
    def __init__(self):
        super(LeNetFeatExtractor, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        return x

class LeNetClassifier(nn.Module):
    def __init__(self):
        super(LeNetClassifier, self).__init__()
        self.fc1 = nn.Linear(16 * 6 * 6, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = torch.flatten(x,1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.feat = LeNetFeatExtractor()
        self.classifer = LeNetClassifier()

    def forward(self, x):
        x = self.feat(x)
        x = self.classifer(x)
        return x

    @torch.jit.export
    def get_length(self, in_len: int) -> int:
        return in_len

[narendasan]

I see. Maybe we could do something like

trtorch.compile(mod, ..., extra_methods=["get_length"])

@peri044 Thoughts?

Are there cases where we wouldnt want to convert forward as well?

[narendasan]

Seems like there are two potential feature possibilities here:

1. The module has methods that do something like get_length that does a simple computation something that is not targeted to be compiled by TRTorch. I see a couple of issues, if these additional methods access members of the module then we need to ensure that the module has a copy of all module attributes. This would increase package size from where we are now. This might require us adding a compress package argument or something to tell TRTorch to not worry about replicating members. The other issue is if these methods are used to modify the state of the module pre-inference. This would cause unexpected behavior since the forward function is frozen using attributes pre-mutation. In theory refitting would solve this but we would need to figure out a way to watch attributes for changes and start refitting if so. This may not be something feasible. We might need to look at FreezeModule preserved attributes.

2. The module has more than one method that needs to be compiled to TRTorch. I think this is something pretty straight forward to implement as long as methods to compile follow rules of TRTorch supported methods i.e pure functions that take the form f(Tensor[]) -> (Tensor[])

[lippman1125]

1. We might need to look at FreezeModule preserved attributes

Yes, I used to make a try by using FreezeModule preserved attributes, but failed.

[narendasan]

It wont be sufficient by FreezeModule alone because TRTorch does not modify the source module in place, it returns a new module with the engine inside.

[github-actions]

This issue has not seen activity for 90 days, Remove stale label or comment or this will be closed in 10 days

[github-actions]

This issue has not seen activity for 90 days, Remove stale label or comment or this will be closed in 10 days