❓ [Question] How do you convert part of the model to TRT?

[MarekPokropinski]

❓ Question

Is it possible to convert only part of the model to TRT. I have model that cannot be directly converted to trt because it uses custom classes. I wanted to convert only modules that can be converted but as I tried it torch cannot save it.

What you have already tried

I tried the following:

import torch.nn
import torch_tensorrt


class MySubmodule(torch.nn.Module):
    def __init__(self):
        super(MySubmodule, self).__init__()
        self.layer = torch.nn.Linear(10, 10)

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


class MyMod(torch.nn.Module):
    def __init__(self):
        super(MyMod, self).__init__()
        self.submod = MySubmodule()
        self.submod = torch_tensorrt.compile(self.submod, inputs=[
            torch_tensorrt.Input(shape=(1, 10))
        ])

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


if __name__ == "__main__":
    model = MyMod()
    scripted = torch.jit.script(model)
    scripted(torch.zeros(1, 10).cuda())
    scripted.save("test.pt")

But it raises exception: RuntimeError: method.qualname() == QualifiedName(selfClass->name()->qualifiedName(), methodName)INTERNAL ASSERT FAILED at "../torch/csrc/jit/serialization/python_print.cpp":1105, please report a bug to PyTorch.

Environment

• PyTorch Version (e.g., 1.0): 1.10.2
• CPU Architecture: intel
• OS (e.g., Linux): linux
• How you installed PyTorch (conda, pip, libtorch, source): pip
• Build command you used (if compiling from source):
• Are you using local sources or building from archives: from archives
• Python version: 3.8
• CUDA version: 11.3
• GPU models and configuration: rtx 3090
• Any other relevant information:

[narendasan]

You have a few options. 1. You can do what you did, but not sure how this how this would effect serialization since you have non PyTorch code in your torchscript. I would like to know what exactly produced this error in PyTorch, but can't look at that right now. 2. You pass the full module (MyMod) to torch_tensorrt.compile and setting MySubmodule as a torch_executed_module. This will tell TorchTRT to represent the submodule in terms of PyTorch operations instead of coalesced into a tensorrt engine. 3. You can also do this at the operator level with torch_executed_ops as well which will do the same for the single op.