Na conversão de modelos ONNX para TensorRT, a definição de formas dinâmicas requer a especificação de três valores: mínimo, ótimo e máximo. Para simplificar procsesos de teste, foi desenvolvido um script que automatiza essa configruação, atribuindo valores padrão a dimensões desconhecidas, enquanto permite ajustes manuais.

A função abaixo analisa um modelo ONNX, identifica automaticamenet as formas dinâmicas e aplica configurações padrão (mínimo 2, ótimo 256, máximo 512) para dimensões variáveis, facilitando a geração de engines TensorRT.

import onnx
import tensorrt as trt
import os

def gerar_engine_trt(caminho_onnx, destino_engine, dtype_trt=trt.DataType.HALF, batch=1, modo_silencioso=False, shapes_dinamicos={}, formas_estaticas={}, limite_memoria=(1 << 30)):
    """Gera um engine TensorRT a partir de um arquivo ONNX para inferência"""

    logger = trt.Logger(trt.Logger.VERBOSE)
    flags_rede = [] if trt.__version__[0] < '7' else [1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)]
    with trt.Builder(logger) as construtor, construtor.create_network(*flags_rede) as rede_trt, trt.OnnxParser(rede_trt, logger) as parser:
        construtor.max_batch_size = batch
        configuracao: trt.IBuilderConfig = construtor.create_builder_config()
        configuracao.max_workspace_size = limite_memoria
        if dtype_trt == trt.DataType.HALF:
            configuracao.set_flag(trt.BuilderFlag.FP16)

        if not os.path.exists(caminho_onnx):
            print(f'Arquivo ONNX {caminho_onnx} inexistente.')
            exit(0)

        print(f'Carregando modelo ONNX de {caminho_onnx}...')
        with open(caminho_onnx, 'rb') as modelo:
            if not parser.parse(modelo.read()):
                print('Erro na análise do arquivo ONNX.')
                for i in range(parser.num_errors):
                    print(parser.get_error(i))
                return None

        print('Análise concluída')
        if not modo_silencioso:
            print(f'Construindo engine a partir de {caminho_onnx}...')

        mapeamento_shapes = {}
        modelo_onnx = onnx.load(caminho_onnx)
        for idx, entrada in enumerate(modelo_onnx.graph.input):
            lista_dimensoes = []
            eh_dinamico = False
            for dimensao in entrada.type.tensor_type.shape.dim:
                lista_dimensoes.append(dimensao.dim_value)
                if dimensao.dim_param or dimensao.dim_value <= 0:
                    eh_dinamico = True
            if eh_dinamico:
                minimo = [(v if v > 0 else 2) for v in lista_dimensoes]
                otimo = [(v if v > 0 else 256) for v in lista_dimensoes]
                maximo = [(v if v > 0 else 512) for v in lista_dimensoes]
                mapeamento_shapes[entrada.name] = [minimo, otimo, maximo]

        for nome, valor in shapes_dinamicos.items():
            mapeamento_shapes[nome] = valor

        perfil_otimizacao = construtor.create_optimization_profile()
        if mapeamento_shapes:
            for nome_tensor, (forma_min, forma_otima, forma_max) in mapeamento_shapes.items():
                perfil_otimizacao.set_shape(nome_tensor, forma_min, forma_otima, forma_max)

        for nome_estatico, (val_min, val_otimo, val_max) in formas_estaticas.items():
            perfil_otimizacao.set_shape_input(nome_estatico, val_min, val_otimo, val_max)

        configuracao.add_optimization_profile(perfil_otimizacao)
        engine = construtor.build_engine(rede_trt, configuracao)
        dados_engine = engine.serialize()
        with open(destino_engine, 'wb') as arquivo_saida:
            arquivo_saida.write(dados_engine)

Exemplo de aplicação: a configuração pode ser manual ou utilizar os valores predefinidos para testes.

nome_do_projeto = "meu_projeto"
gerar_engine_trt(f"onnx/{nome_do_projeto}/encoder.onnx", f"onnx/{nome_do_projeto}/encoder.trt",
    shapes_dinamicos={
        "seq_referencia": [(1, 1), (1, 256), (1, 512)],
        "seq_texto": [(1, 1), (1, 256), (1, 512)],
        "bert_ref": [(1024, 1), (1024, 256), (1024, 512)],
        "bert_texto": [(1024, 1), (1024, 256), (1024, 512)],
        "ssl_saida": [(1, 768, 1), (1, 768, 256), (1, 768, 512)],
    },
    formas_estaticas={
        "comprimento_saida": [(165,), (265,), (1165,)],
        "comprimento_chave_valor": [(227,), (327,), (1227,)]
    }
)

gerar_engine_trt(f"onnx/{nome_do_projeto}/decodificador.onnx", f"onnx/{nome_do_projeto}/decodificador.trt")