❓ [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.

[ashafaei]

@narendasan This solution is the opposite of what's needed here. We want to preserve the high-level MyMod as the standard torch script, but convert a submodule to TensorRT. Solutions 2, 3 specify what high-level object should stay as torch script, while a better structure may be to say which low-level objects should be compiled to TensorRT instead.

The original problem is not this though. We can manually compile a submodule as TensorRT, and then trace the high-level object using Torch, and it'll work just well. The issue is when we try to serialize the high-level TS that now partially contains TensorRT. The above code perfectly shows the context and the issue.

I've been struggling with this for the past few hours. the best solution I have so far is to separately save/load the TRT submodules and assemble the final model in runtime.

[MatthieuToulemont]

@narendasan

Let's say we have the following set up :

class Pipeline(torch.nn.Module):
    def __init__(self, preprocessing: PreprocessingModule, model, postprocessing: PostprocessingModule) -> None:
        super().__init__()

        self.preprocessing = preprocessing
        self.model = model
        self.postprocessing = postprocessing
    
    def forward(self, image):
        preprocessed_image = self.preprocessing(image)
        model_prediction = self.model(preprocessing_image)
        return self.postprocessing(image, model_prediction)

Where self.preprocessing and self.postprocessing are modules that are JIT scriptable but not TRT compatible.
On the other hand self.model is not JIT scriptable but JIT traceable and TRT compatible.

Being able to compile self.model with Torch-Tensor RT and then add it to the pipeline which is overall JIT scriptable enables us to create checkpoints we can easily deploy to production with Nvidia Triton Inference Server for instance.

Needing to save the TRT compiled model in a separate checkpoint makes the deployment process longer as we need to deploy 3 models instead of one (in the set up described above).

[ekolodin]

Did you manage with this problem? I had the same issue :(

[t90tank]

same issue :(