The purpose of this example is to demonstrate the overall flow of lowering a PyTorch model to TensorRT via FX with existing FX based tooling. The general lowering flow would be like:
- Use splitter to split the model if there're ops in the model that we don't want to lower to TensorRT for some reasons like the ops are not supported in TensorRT or running them on other backends provides better performance.
- Lower the model (or part of the model if splitter is used) to TensorRT via fx path. If we know the model is fully supported by fx path (without op unsupported) then we can skip the splitter.
In [1]:
import torch
import torch.fx
import torch.nn as nn
from torch_tensorrt.fx.utils import LowerPrecision
import torch_tensorrt.fx.tracer.acc_tracer.acc_tracer as acc_tracer
from torch_tensorrt.fx import InputTensorSpec, TRTInterpreter, TRTModule
from torch_tensorrt.fx.tools.trt_splitter import TRTSplitter
In [2]:
class Model(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 10)
self.relu = nn.ReLU()
def forward(self, x):
x = self.linear(x)
x = self.relu(x)
x = torch.linalg.norm(x, ord=2, dim=1)
x = self.relu(x)
return x
inputs = [torch.randn((1, 10), device=torch.device('cuda'))]
model = Model().cuda().eval()
acc_tracer is a custom fx tracer that maps nodes whose targets are PyTorch operators to acc ops.
In [3]:
traced = acc_tracer.trace(model, inputs)
Splitter will split the model into several submodules. The name of submodules will
be either run_on_acc_{} or run_on_gpu_{}. Submodules named run_on_acc_{} can
be fully lowered to TensorRT via fx2trt while submodules named run_on_gpu_{} has
unsupported ops and can't be lowered by fx2trt. We can still run run_on_gpu_{}
submodules on GPU if ops there have cuda implementation.
In [4]:
splitter = TRTSplitter(traced, inputs)
Preview functionality allows us to see what are the supported ops and unsupported ops. We can optionally the dot graph which will color supported ops and unsupported ops differently.
In [5]:
splitter.node_support_preview()
Supported node types in the model:
acc_ops.linear: ((), {'input': torch.float32, 'weight': torch.float32, 'bias': torch.float32})
acc_ops.relu: ((), {'input': torch.float32})
Unsupported node types in the model:
acc_ops.linalg_norm: ((), {'input': torch.float32})
Out[5]:
"\nSupported node types in the model:\nacc_ops.linear: ((), {'input': torch.float32, 'weight': torch.float32, 'bias': torch.float32})\nacc_ops.relu: ((), {'input': torch.float32})\n\nUnsupported node types in the model:\nacc_ops.linalg_norm: ((), {'input': torch.float32})\n"
After split, there are three submodules, _run_on_acc_0 and _run_on_gpu_1.
In [6]:
split_mod = splitter()
print(split_mod.graph)
Got 2 acc subgraphs and 1 non-acc subgraphs
graph():
%x : [#users=1] = placeholder[target=x]
%_run_on_acc_0 : [#users=1] = call_module[target=_run_on_acc_0](args = (%x,), kwargs = {})
%_run_on_gpu_1 : [#users=1] = call_module[target=_run_on_gpu_1](args = (%_run_on_acc_0,), kwargs = {})
%_run_on_acc_2 : [#users=1] = call_module[target=_run_on_acc_2](args = (%_run_on_gpu_1,), kwargs = {})
return _run_on_acc_2
In [7]:
print(split_mod._run_on_acc_0.graph)
print(split_mod._run_on_gpu_1.graph)
print(split_mod._run_on_acc_2.graph)
graph():
%x : [#users=1] = placeholder[target=x]
%linear_weight : [#users=1] = get_attr[target=linear.weight]
%linear_bias : [#users=1] = get_attr[target=linear.bias]
%linear_1 : [#users=1] = call_function[target=torch_tensorrt.fx.tracer.acc_tracer.acc_ops.linear](args = (), kwargs = {input: %x, weight: %linear_weight, bias: %linear_bias})
%relu_2 : [#users=1] = call_function[target=torch_tensorrt.fx.tracer.acc_tracer.acc_ops.relu](args = (), kwargs = {input: %linear_1, inplace: False})
return relu_2
graph():
%relu_2 : [#users=1] = placeholder[target=relu_2]
%linalg_norm_1 : [#users=1] = call_function[target=torch_tensorrt.fx.tracer.acc_tracer.acc_ops.linalg_norm](args = (), kwargs = {input: %relu_2, ord: 2, dim: 1, keepdim: False})
return linalg_norm_1
graph():
%linalg_norm_1 : [#users=1] = placeholder[target=linalg_norm_1]
%relu_3 : [#users=1] = call_function[target=torch_tensorrt.fx.tracer.acc_tracer.acc_ops.relu](args = (), kwargs = {input: %linalg_norm_1, inplace: False})
return relu_3
The split_mod contains the child modules supported by TRT or eager gpu. We can iterate them to transform the module into TRT engine.
In [8]:
def get_submod_inputs(mod, submod, inputs):
acc_inputs = None
def get_input(self, inputs):
nonlocal acc_inputs
acc_inputs = inputs
handle = submod.register_forward_pre_hook(get_input)
mod(*inputs)
handle.remove()
return acc_inputs
# Since the model is splitted into three segments. We need to lower each TRT eligible segment.
# If we know the model can be fully lowered, we can skip the splitter part.
for name, _ in split_mod.named_children():
if "_run_on_acc" in name:
submod = getattr(split_mod, name)
# Get submodule inputs for fx2trt
acc_inputs = get_submod_inputs(split_mod, submod, inputs)
# fx2trt replacement
interp = TRTInterpreter(
submod,
InputTensorSpec.from_tensors(acc_inputs),
explicit_batch_dimension=True,
)
r = interp.run(lower_precision=LowerPrecision.FP32)
trt_mod = TRTModule(*r)
setattr(split_mod, name, trt_mod)
lowered_model_output = split_mod(*inputs)
I0627 150503.073 fx2trt.py:190] Run Module elapsed time: 0:00:00.014965
I0627 150504.996 fx2trt.py:241] Build TRT engine elapsed time: 0:00:01.922029
I0627 150505.026 fx2trt.py:190] Run Module elapsed time: 0:00:00.000302
I0627 150509.953 fx2trt.py:241] Build TRT engine elapsed time: 0:00:04.925192
Model can be saved by torch.save and loaded with torch.load. Then we can compare the results with eager mode inference.
In [9]:
torch.save(split_mod, "trt.pt")
reload_trt_mod = torch.load("trt.pt")
reload_model_output = reload_trt_mod(*inputs)
# Make sure the results match
regular_model_output = model(*inputs)
torch.testing.assert_close(
reload_model_output, regular_model_output, atol=3e-3, rtol=1e-2
)