Refit Engine given PyTorch `state_dict`

[nrakltx]

I have a TensorRT model and I would like to refit some of its weights.

I created this model by exporting a PyTorch model to ONNX and from there compiling to TensorRT.
It would be impractical to re-compile (+re-export) the model to ONNX and TRT as it is big and this needs to happen many times.

Given my modified torch state_dict, I would like to be automatically able to refit the TRT engine's weights accordingly.

From my research, it seems exporting torch to ONNX doesn't provide a 1:1 mapping of weight names (or shapes), and even less so ONNX -> TRT. If no automatic method is found, it becomes far harder to refit, as a manual mapping of each of the (many) weights of the model would need to be computed by a human.

Is there some way this could still be done?
If not, could an ONNX -> TensorRT mapping be deduced somehow at least?

Thank you very much in advance,

N

[zerollzeng]

the name mapping from Pytorch to ONNX is not what TRT can control so I don't know how to solve this :(

For refit an engine, does https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#refitting-engine-c help?

[nrakltx]

Thanks for the prompt response @zerollzeng!
I was already looking at the Refitter API, but the problem is that it requires names of weights in the final engine file namespace, and not the ONNX namespace.
Since I am not able to (humanly, manually) follow the ONNX -> TRT mapping, I would like to be able to talk with the API in ONNX terms.
This is what I tried asking about originally - do you have a direction for me to follow?

[zerollzeng]

Sorry I don't know much about the details of refit. but IIRC the output engine layers name should be the combination of the original onnx layers names. like

[12/27/2022-15:06:39] [V] [TRT] Engine Layer Information:
Layer(Reformat): Reformatting CopyNode for Input Tensor 0 to node_of_gpu_0/conv1_1 + node_of_gpu_0/res_conv1_bn_1 + node_of_gpu_0/res_conv1_bn_2 + node_of_gpu_0/pool1_1, Tactic: 0x00000000000003e8, gpu_0/data_0 (Float[1,3,224,224]) -> Reformatted Input Tensor 0 to node_of_gpu_0/conv1_1 + node_of_gpu_0/res_conv1_bn_1 + node_of_gpu_0/res_conv1_bn_2 + node_of_gpu_0/pool1_1 (Int8[1,3,224,224])
Layer(CaskConvActPool): node_of_gpu_0/conv1_1 + node_of_gpu_0/res_conv1_bn_1 + node_of_gpu_0/res_conv1_bn_2 + node_of_gpu_0/pool1_1, Tactic: 0xb71c75095873646c, Reformatted Input Tensor 0 to node_of_gpu_0/conv1_1 + node_of_gpu_0/res_conv1_bn_1 + node_of_gpu_0/res_conv1_bn_2 + node_of_gpu_0/pool1_1 (Int8[1,3,224,224]) -> gpu_0/pool1_1 (Int8[1,64:32,56,56])
Layer(CaskConvolution): node_of_gpu_0/res2_0_branch2a_1 + node_of_gpu_0/res2_0_branch2a_bn_1 + node_of_gpu_0/res2_0_branch2a_bn_2, Tactic: 0x4f35593c356e2e7e, gpu_0/pool1_1 (Int8[1,64:32,56,56]) -> gpu_0/res2_0_branch2a_bn_2 (Int8[1,64:32,56,56])
Layer(CaskConvolution): node_of_gpu_0/res2_0_branch2b_1 + node_of_gpu_0/res2_0_branch2b_bn_1 + node_of_gpu_0/res2_0_branch2b_bn_2, Tactic: 0xd3d41ef6de22d9b6, gpu_0/res2_0_branch2a_bn_2 (Int8[1,64:32,56,56]) -> gpu_0/res2_0_branch2b_bn_2 (Int8[1,64:32,56,56])
Layer(CaskConvolution): node_of_gpu_0/res2_0_branch2c_1 + node_of_gpu_0/res2_0_branch2c_bn_1, Tactic: 0x5e4918ccf433630e, gpu_0/res2_0_branch2b_bn_2 (Int8[1,64:32,56,56]) -> gpu_0/res2_0_branch2c_bn_1 (Int8[1,256:32,56,56])

[zerollzeng]

@ttyio Do you have any suggestions here? Thanks!

[nrakltx]

@zerollzeng that might be interesting - is there any way to parse this logging output into a mapping of the values? (say, for each TRT weight, a dictionary of ONNX weights and where each of them affects the TRT weight).
Alternatively, could you point me to some docs on the format of this logging output?

Thank you so much!

[zerollzeng]

I don't think we have such docs but this is possible.

[ttyio]

@nrakltx , Here is a solution Ideally should work.

let's assume you have 2 onnx, our goal is to use first onnx to generate a refittable engine, and use the second onnx to update weights.

After generate trt engine from the 1st onnx, let's use onnx parser to load the 2nd onnx, and iterator the INetworkDefination. Check all the interested constantNode , convolutionNode and other node that you are interested. You can get a mapping from tensor name to tensor (constant and conv weights/bias). Finally let loop all the refit.getAllWeights, find the tensor from the previous mapping and update the weights.

Hope this helpful.

[ttyio]

closing since no activity for more than 3 weeks, thank you!