Is it possible to convert PyTorch multimodal CVAE to TensorRT?

tlgkyn06: <h2>
<a name="description-1" class="anchor" href="#description-1"></a>Description</h2>
<p>Hi, I’m working on behavioral prediction code which uses multimodal CVAE (link:<a href="https://github.com/StanfordASL/Trajectron-plus-plus" rel="noopener nofollow ugc">GitHub - StanfordASL/Trajectron-plus-plus: Code accompanying the ECCV 2020 paper “Trajectron++: Dynamically-Feasible Trajectory Forecasting With Heterogeneous Data” by Tim Salzmann*, Boris Ivanovic*, Punarjay Chakravarty, and Marco Pavone (* denotes equal contribution).</a>) in examples the common approaches is pytorch to onnx then onnx to tensorrt. But when the model is decode-encode structured and multimodal, I could not find the way how to use torch.onnx.export() function.<br>
pytorch model link: <a href="https://github.com/StanfordASL/Trajectron-plus-plus/tree/eccv2020/experiments/nuScenes/models/int_ee" class="inline-onebox" rel="noopener nofollow ugc">Trajectron-plus-plus/experiments/nuScenes/models/int_ee at eccv2020 · StanfordASL/Trajectron-plus-plus · GitHub</a></p>
<p><strong>torch.onnx.export() function accepts these arguments:</strong></p>
<p>torch.onnx.export(torch_model,               # model being run<br>
x,                         # model input (or a tuple for multiple inputs)<br>
“super_resolution.onnx”,   # where to save the model (can be a file or file-like object)<br>
export_params=True,        # store the trained parameter weights inside the model file<br>
opset_version=10,          # the ONNX version to export the model to<br>
do_constant_folding=True,  # whether to execute constant folding for optimization<br>
input_names = [‘input’],   # the model’s input names<br>
output_names = [‘output’], # the model’s output names<br>
dynamic_axes={‘input’ : {0 : ‘batch_size’},    # variable length axes<br>
‘output’ : {0 : ‘batch_size’}})</p>
<h1>
<a name="according-to-informations-that-i-indicated-below-what-should-input_names-output_names-and-dynamic-axes-be-2" class="anchor" href="#according-to-informations-that-i-indicated-below-what-should-input_names-output_names-and-dynamic-axes-be-2"></a>According to informations that I indicated below, what should input_names, output_names, and dynamic axes be?</h1>
<p>As it is shown at below my input changes according to classification (Vehicle or Pedestrian)</p>
<p><strong>the models encode function is this :</strong><br>
self.node_models_dict[node].encoder_forward(inputs,<br>
inputs_st,<br>
inputs_np,<br>
robot_present_and_future,<br>
maps)<br>
<strong>Example encode inputs are like that:</strong><br>
inputs= {VEHICLE/35521: tensor([[-2.5944e+01, 5.7025e+00, -5.8111e+00, 1.1085e-03, -8.1415e-01,<br>
4.1824e-02, 1.8372e+02, 7.1199e-02]], device=‘cuda:0’)}<br>
inputs_st= {VEHICLE/35521: tensor([[ 0.0000e+00, 0.0000e+00, -3.8740e-01, 7.3899e-05, -2.0354e-01,<br>
1.0456e-02, 5.8481e+01, 7.1199e-02]], device=‘cuda:0’)}<br>
inputs_np= {VEHICLE/35521: array([[-2.59436354e+01, 5.70252158e+00, -5.81107431e+00,<br>
1.10847897e-03, -8.14148775e-01, 4.18236967e-02,<br>
1.83724350e+02, 7.11987019e-02]])}<br>
robot_present_and_future: None<br>
maps: None</p>
<p><strong>The models decode function is this :</strong><br>
model.decoder_forward(num_predicted_timesteps,<br>
num_samples,<br>
robot_present_and_future=robot_present_and_future,<br>
z_mode=z_mode,<br>
gmm_mode=gmm_mode,<br>
full_dist=full_dist,<br>
all_z_sep=all_z_sep)</p>
<p><strong>Example decode inputs arelike that:</strong><br>
node(encoder): VEHICLE/35521<br>
node: VEHICLE/35521<br>
num_predicted_timesteps: 6<br>
num_samples: 1<br>
robot_present_and_future: None<br>
z_mode: False<br>
gmm_mode: False<br>
full_dist: True<br>
all_z_sep: False</p>
<p><strong>Pytorch model_dict was shown below:</strong><br>
name: VEHICLE/node_history_encoder model: LSTM(8, 32, batch_first=True)<br>
name: VEHICLE/node_future_encoder model: LSTM(2, 32, batch_first=True, bidirectional=True)<br>
name: VEHICLE/node_future_encoder/initial_h model: Linear(in_features=8, out_features=32, bias=True)<br>
name: VEHICLE/node_future_encoder/initial_c model: Linear(in_features=8, out_features=32, bias=True)<br>
name: VEHICLE/edge_influence_encoder model: AdditiveAttention(<br>
(w1): Linear(in_features=32, out_features=32, bias=False)<br>
(w2): Linear(in_features=32, out_features=32, bias=False)<br>
(v): Linear(in_features=32, out_features=1, bias=False)<br>
)<br>
name: VEHICLE/p_z_x model: Linear(in_features=64, out_features=32, bias=True)<br>
name: VEHICLE/hx_to_z model: Linear(in_features=32, out_features=25, bias=True)<br>
name: VEHICLE/hxy_to_z model: Linear(in_features=192, out_features=25, bias=True)<br>
name: VEHICLE/decoder/state_action model: Sequential(<br>
(0): Linear(in_features=8, out_features=2, bias=True)<br>
)<br>
name: VEHICLE/decoder/rnn_cell model: GRUCell(91, 128)<br>
name: VEHICLE/decoder/initial_h model: Linear(in_features=89, out_features=128, bias=True)<br>
name: VEHICLE/decoder/proj_to_GMM_log_pis model: Linear(in_features=128, out_features=1, bias=True)<br>
name: VEHICLE/decoder/proj_to_GMM_mus model: Linear(in_features=128, out_features=2, bias=True)<br>
name: VEHICLE/decoder/proj_to_GMM_log_sigmas model: Linear(in_features=128, out_features=2, bias=True)<br>
name: VEHICLE/decoder/proj_to_GMM_corrs model: Linear(in_features=128, out_features=1, bias=True)<br>
name: VEHICLE-&gt;VEHICLE/edge_encoder model: LSTM(16, 32, batch_first=True)<br>
name: VEHICLE-&gt;PEDESTRIAN/edge_encoder model: LSTM(14, 32, batch_first=True)<br>
name: VEHICLE/unicycle_initializer model: Linear(in_features=65, out_features=1, bias=True)<br>
name: PEDESTRIAN/node_history_encoder model: LSTM(6, 32, batch_first=True)<br>
name: PEDESTRIAN/node_future_encoder model: LSTM(2, 32, batch_first=True, bidirectional=True)<br>
name: PEDESTRIAN/node_future_encoder/initial_h model: Linear(in_features=6, out_features=32, bias=True)<br>
name: PEDESTRIAN/node_future_encoder/initial_c model: Linear(in_features=6, out_features=32, bias=True)<br>
name: PEDESTRIAN/edge_influence_encoder model: AdditiveAttention(<br>
(w1): Linear(in_features=32, out_features=32, bias=False)<br>
(w2): Linear(in_features=32, out_features=32, bias=False)<br>
(v): Linear(in_features=32, out_features=1, bias=False)<br>
)<br>
name: PEDESTRIAN/p_z_x model: Linear(in_features=64, out_features=32, bias=True)<br>
name: PEDESTRIAN/hx_to_z model: Linear(in_features=32, out_features=25, bias=True)<br>
name: PEDESTRIAN/hxy_to_z model: Linear(in_features=192, out_features=25, bias=True)<br>
name: PEDESTRIAN/decoder/state_action model: Sequential(<br>
(0): Linear(in_features=6, out_features=2, bias=True)<br>
)<br>
name: PEDESTRIAN/decoder/rnn_cell model: GRUCell(91, 128)<br>
name: PEDESTRIAN/decoder/initial_h model: Linear(in_features=89, out_features=128, bias=True)<br>
name: PEDESTRIAN/decoder/proj_to_GMM_log_pis model: Linear(in_features=128, out_features=1, bias=True)<br>
name: PEDESTRIAN/decoder/proj_to_GMM_mus model: Linear(in_features=128, out_features=2, bias=True)<br>
name: PEDESTRIAN/decoder/proj_to_GMM_log_sigmas model: Linear(in_features=128, out_features=2, bias=True)<br>
name: PEDESTRIAN/decoder/proj_to_GMM_corrs model: Linear(in_features=128, out_features=1, bias=True)<br>
name: PEDESTRIAN-&gt;VEHICLE/edge_encoder model: LSTM(14, 32, batch_first=True)<br>
name: PEDESTRIAN-&gt;PEDESTRIAN/edge_encoder model: LSTM(12, 32, batch_first=True)</p>