FastPitch TensorRT infer

[yuananf]

I can convert FastPitch model to tensorrt, but I can't do the infer, because the output shape of spect can not be determined.

Anyone can help?

NVIDIA/TensorRT#4098

[zhenyih]

AI generated solution by traversing NeMo source. Please verify.

To solve the issue with FastPitch TensorRT inference where the output shape of 'spect' cannot be determined, you need to explicitly define dynamic output dimensions when converting to TensorRT.

The problem occurs because the spectrogram length in FastPitch is dynamically determined during inference based on the predicted durations of each input token. TensorRT requires all tensor shapes to be known at build time or to have explicitly defined dynamic dimensions.

Here's how to fix it:

1. When exporting to ONNX, make sure to specify dynamic axes for both inputs and outputs:

dynamic_axes = {
    'text': {0: 'batch_size', 1: 'text_length'},
    'spect': {0: 'batch_size', 2: 'spec_length'}  # The spectrogram length is dynamic
}
torch.onnx.export(model, args, "fastpitch.onnx", dynamic_axes=dynamic_axes)

2. When building the TensorRT engine, define both input and output profiles with minimum, optimal, and maximum dimensions:

# For input profiles
p.add("text", min=(1, 1), opt=(1, 50), max=(1, 150))
p.add("pitch", min=(1, 1), opt=(1, 50), max=(1, 150))

# For output profiles - specify the range for the spectrogram length
p.add("spect", min=(1, 80, 10), opt=(1, 80, 200), max=(1, 80, 600))

The key is to provide TensorRT with information about the possible range of your output spectrogram length. The exact values will depend on your specific use case, but you should set the maximum length generously to accommodate longer outputs.

If you're using the NeMo export API, you may need to modify the export configuration to include these dynamic output dimensions.