✨[Feature] Add detailed layer name to the NVTX markers in torch-TRT #2069
Description
Is your feature request related to a problem? Please describe.
I use the nsys profiler regularly to study network performance. One of the important things I need to know is which unique layers map to specific kernels. There are currently nvtx markers generated by TRT which show some of the layer specific information, but the layer name is not useful. The layer name is currently unique eg conv_20 or mul_40 but it doesn't provide enough context to identify exactly where in the network graph the operator is being executed. PyTorch provides a mechanism to record the detailed unique name of each nn.module instance. This name can be pre-fixed to the current operator name to get a unique name which tells you exactly where the node is found in the network architecture.
Example from ResNet
/layer1/0/conv1/convolution_1
This shows that the convolution is from the 1st bottleneck block in resnet (/layer1). Each bottleneck is composed of 2 instances of the same type, implemented using a module list, this example is instance 0 (layer1/0) . Then it is the 1st conv (conv1). Lastly the operator type is a convolution (convolution_1). With this information it is seldom necessary to inspect the network source code to understand the network architecture.
Describe the solution you'd like
Torch-trt is just a pass through for this case. The module name is created by PyTorch and TensorRT already has NVTX markers. The only thing needed is for torch-TRT to capture the module name and pass it to TRT. This layer name is currently passed to TRT here in run_node() from fx2trt. The change is just to fetch the module name from pytorch and prefix it to the node name before passing it to TRT.
Describe alternatives you've considered
There aren't any good alternatives that I'm aware of.
Additional context
Here's an example of what the detailed layer name looks like when it is included in the TRT NVTX markers. I have already implemented this feature in my local copy of torch-TRT.
