Neste guia, abordaremos a instalação e configuração do NVIDIA TensorRT no sistema operacional Ubuntu 20.04, focando na sua aplicação para acelerar a inferência de modelos de detecção de objetos YOLOv5. O TensorRT é um SDK de alto desempenho para inferência de aprendizado profundo que oferece otimizações significativas para redes neurais, resultando em latência reduzida e maior throughput.
- Download do Kit de Desenvolvimento TensorRT
Para iniciar, é necessário baixar o pacote do TensorRT diretamente do portal de desenovlvedores da NVIDIA. Certifique-se de escolher a versão compatível com sua versão do CUDA e cuDNN.
Após o download, mova o arquivo compactado para um diretório de sua preferência (por exemplo, seu diretório home) e descompacte-o:
tar -zxvf TensorRT-X.Y.Z.Linux.x86_64-gnu.cuda-W.V.cudnnP.Q.tar.gz
Substitua TensorRT-X.Y.Z.Linux.x86_64-gnu.cuda-W.V.cudnnP.Q.tar.gz pelo nome exato do arquivo que você baixou.
- Configuração das Variáveis de Ambiente
Para que o sistema reconheça as bibliotecas do TensorRT, é preciso adicionar o caminho para o diretório lib do TensorRT à variável de ambiente LD_LIBRARY_PATH.
Edite o arquivo de configuração do seu shell (geralmente .bashrc para Bash):
nano ~/.bashrc
Adicione a seguinte linha ao final do arquivo, ajustando o caminho para o diretório onde você descompactou o TensorRT (neste exemplo, /home/seu_usuario/TensorRT-8.2.1.8):
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/seu_usuario/TensorRT-8.2.1.8/lib
Salve o arquivo e saia do editor. Em seguida, recarregue as variáveis de ambiente para aplicar as mudanças:
source ~/.bashrc
Com essas etapas, a instalação básica do TensorRT está concluída.
- Validação Funcional com YOLOv5
Para testar a instalação do TensorRT, utilizaremos um projeto de exemplo que integra o YOLOv5 com o TensorRT. É crucial usar versões compatíveis do modelo YOLOv5 e do projeto de aceleração.
3.1. Obtenção do Projeto e Modelos
Baixe o repositório tensorrtx, que contém os scripts e implementações para converter e executar modelos YOLOv5 com TensorRT:
- Repositório
tensorrtx: https://gitcode.net/mirrors/wang-xinyu/tensorrtx
Em seguida, baixe o modelo YOLOv5 (arquivo .pt) correspondente à versão do tensorrtx que você está utilizando. É fundamental que as versões do modelo e do código sejam compatíveis.
- Modelos YOLOv5: https://github.com/ultralytics/yolov5/releases
Após baixar o repositório tensorrtx e o modelo YOLOv5 (por exemplo, yolov5s.pt), navegue até o diretório yolov5 dentro do repositório tensorrtx.
3.2. Geração do Arquivo de Pesos (.wts)
Dentro do diretório yolov5 (do tensorrtx), o arquivo gen_wts.py é utilizado para extrair os pesos do modelo PyTorch (.pt) e convertê-los para o formato .wts, que é necessário para a construção do motor TensorRT. Copie gen_wts.py para o diretório raiz do seu clone do YOLOv5 (onde está o .pt) ou execute-o a partir do diretório tensorrtx/yolov5, garantindo que o .pt esteja acessível.
Execute o script Python para gerar o arquivo de pesos. Certifique-se de que o arquivo yolov5s.pt esteja no mesmo diretório ou forneça o caminho completo.
python3 gen_wts.py --w yolov5s.pt
3.3. Compilação do Projeto TensorRTx
Agora, compile o projeto C++ do TensorRTx. Ainda dentro do diretório yolov5 do tensorrtx:
mkdir build && cd build
cmake ..
make
Após a compilação bem-sucedida, você terá um executável chamado yolov5 no diretório build. Mova o arquivo .wts gerado na etapa anterior para este mesmo diretório build.
3.4. Criação e Execução do Motor de Inferência
Primerio, crie o arquivo de motor do TensorRT (.engine) a partir dos pesos .wts:
./yolov5 -s yolov5s.wts yolov5s.engine s
Se você enfrentar erros de permissão, tente prefixar os comandos com sudo, mas use-o com cautela.
Para testar a inferência em um conjunto de imagens (assumindo que você tenha um diretório samples com imagens no diretório tensorrtx/yolov5):
./yolov5 -d yolov5s.engine ../samples
3.5. Habilitando a Detecção por Câmera
Para realizar a inferência em tempo real usando uma câmera, é necessário modificar o código-fonte C++ do yolov5.cpp. O seguinte trecho de código demonstra a configuração para captura de vídeo, pré-processamento de frames e execução da inferência com TensorRT.
Substitua o conteúdo do seu arquivo yolov5.cpp (localizado no diretório tensorrtx/yolov5) pelo código abaixo:
#include <iostream>
#include <chrono>
#include <string>
#include <vector>
#include <fstream>
#include <cassert>
#include <map>
#include <cmath>
#include "cuda_utils.h" // Assumindo utilitários CUDA
#include "logging.h" // Assumindo logger do TensorRT
#include "common.hpp" // Assumindo definições comuns do YOLO
#include "utils.h" // Assumindo utilitários de imagem e NMS
#include "calibrator.h" // Para INT8
#include <opencv2/opencv.hpp> // Para processamento de imagem e vídeo
#define ENABLE_FP16_MODE // Defina ENABLE_INT8_MODE, ENABLE_FP16_MODE ou ENABLE_FP32_MODE
#define GPU_DEVICE_ID 0 // ID da GPU a ser utilizada
#define NMS_THRESHOLD_VAL 0.4
#define CONFIDENCE_THRESHOLD_VAL 0.5
#define INFERENCE_BATCH_SIZE 1
// Parâmetros da rede e dos tensores de entrada/saída
static const int MODEL_INPUT_HEIGHT = Yolo::INPUT_H;
static const int MODEL_INPUT_WIDTH = Yolo::INPUT_W;
static const int NUM_CLASSES = Yolo::CLASS_NUM;
static const int DETECTION_OUTPUT_SIZE = Yolo::MAX_OUTPUT_BBOX_COUNT * sizeof(Yolo::Detection) / sizeof(float) + 1; // Saída máxima de caixas
const char* INPUT_TENSOR_NAME = "data";
const char* OUTPUT_TENSOR_NAME = "prob";
static Logger trt_logger; // Objeto de logging para TensorRT
// Classes de detecção (exemplo para COCO)
const char* object_classes[] = {
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
"hair drier", "toothbrush"
};
// Funções utilitárias para cálculo de dimensões de camadas
static int calculate_padded_width(int x, float width_factor, int alignment_divisor = 8) {
return (int)(std::ceil(x * width_factor / alignment_divisor) * alignment_divisor);
}
static int calculate_padded_depth(int x, float depth_factor) {
if (x == 1) return 1;
return (int)(std::round(x * depth_factor) > 1 ? std::round(x * depth_factor) : 1);
}
// Constrói o motor do TensorRT para o modelo YOLOv5
nvinfer1::ICudaEngine* build_yolov5_engine(
unsigned int batch_size_limit,
nvinfer1::IBuilder* trt_builder,
nvinfer1::IBuilderConfig* builder_config,
nvinfer1::DataType data_type,
float& depth_factor,
float& width_factor,
std::string& weights_file_path)
{
nvinfer1::INetworkDefinition* model_network = trt_builder->createNetworkV2(0U);
assert(model_network);
// Cria o tensor de entrada com formato {3, MODEL_INPUT_HEIGHT, MODEL_INPUT_WIDTH}
nvinfer1::ITensor* input_tensor = model_network->addInput(INPUT_TENSOR_NAME, data_type, nvinfer1::Dims3{ 3, MODEL_INPUT_HEIGHT, MODEL_INPUT_WIDTH });
assert(input_tensor);
std::map<:string weights=""> weights_map = loadWeights(weights_file_path);
/* ------ Backbone do YOLOv5 ------ */
auto focus_layer_0 = focus(model_network, weights_map, *input_tensor, 3, calculate_padded_width(64, width_factor), 3, "model.0");
auto conv_block_1 = convBlock(model_network, weights_map, *focus_layer_0->getOutput(0), calculate_padded_width(128, width_factor), 3, 2, 1, "model.1");
auto c3_block_2 = C3(model_network, weights_map, *conv_block_1->getOutput(0), calculate_padded_width(128, width_factor), calculate_padded_width(128, width_factor), calculate_padded_depth(3, depth_factor), true, 1, 0.5, "model.2");
auto conv_block_3 = convBlock(model_network, weights_map, *c3_block_2->getOutput(0), calculate_padded_width(256, width_factor), 3, 2, 1, "model.3");
auto c3_block_4 = C3(model_network, weights_map, *conv_block_3->getOutput(0), calculate_padded_width(256, width_factor), calculate_padded_width(256, width_factor), calculate_padded_depth(9, depth_factor), true, 1, 0.5, "model.4");
auto conv_block_5 = convBlock(model_network, weights_map, *c3_block_4->getOutput(0), calculate_padded_width(512, width_factor), 3, 2, 1, "model.5");
auto c3_block_6 = C3(model_network, weights_map, *conv_block_5->getOutput(0), calculate_padded_width(512, width_factor), calculate_padded_width(512, width_factor), calculate_padded_depth(9, depth_factor), true, 1, 0.5, "model.6");
auto conv_block_7 = convBlock(model_network, weights_map, *c3_block_6->getOutput(0), calculate_padded_width(1024, width_factor), 3, 2, 1, "model.7");
auto spp_block_8 = SPP(model_network, weights_map, *conv_block_7->getOutput(0), calculate_padded_width(1024, width_factor), calculate_padded_width(1024, width_factor), 5, 9, 13, "model.8");
/* ------ Head do YOLOv5 ------ */
auto c3_block_9 = C3(model_network, weights_map, *spp_block_8->getOutput(0), calculate_padded_width(1024, width_factor), calculate_padded_width(1024, width_factor), calculate_padded_depth(3, depth_factor), false, 1, 0.5, "model.9");
auto conv_block_10 = convBlock(model_network, weights_map, *c3_block_9->getOutput(0), calculate_padded_width(512, width_factor), 1, 1, 1, "model.10");
auto upsample_layer_11 = model_network->addResize(*conv_block_10->getOutput(0));
assert(upsample_layer_11);
upsample_layer_11->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
upsample_layer_11->setOutputDimensions(c3_block_6->getOutput(0)->getDimensions());
nvinfer1::ITensor* concat_inputs_12[] = { upsample_layer_11->getOutput(0), c3_block_6->getOutput(0) };
auto concat_layer_12 = model_network->addConcatenation(concat_inputs_12, 2);
auto c3_block_13 = C3(model_network, weights_map, *concat_layer_12->getOutput(0), calculate_padded_width(1024, width_factor), calculate_padded_width(512, width_factor), calculate_padded_depth(3, depth_factor), false, 1, 0.5, "model.13");
auto conv_block_14 = convBlock(model_network, weights_map, *c3_block_13->getOutput(0), calculate_padded_width(256, width_factor), 1, 1, 1, "model.14");
auto upsample_layer_15 = model_network->addResize(*conv_block_14->getOutput(0));
assert(upsample_layer_15);
upsample_layer_15->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
upsample_layer_15->setOutputDimensions(c3_block_4->getOutput(0)->getDimensions());
nvinfer1::ITensor* concat_inputs_16[] = { upsample_layer_15->getOutput(0), c3_block_4->getOutput(0) };
auto concat_layer_16 = model_network->addConcatenation(concat_inputs_16, 2);
auto c3_block_17 = C3(model_network, weights_map, *concat_layer_16->getOutput(0), calculate_padded_width(512, width_factor), calculate_padded_width(256, width_factor), calculate_padded_depth(3, depth_factor), false, 1, 0.5, "model.17");
// Camada YOLO 0
nvinfer1::IConvolutionLayer* detection_layer_0 = model_network->addConvolutionNd(*c3_block_17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), nvinfer1::DimsHW{ 1, 1 }, weights_map["model.24.m.0.weight"], weights_map["model.24.m.0.bias"]);
auto conv_block_18 = convBlock(model_network, weights_map, *c3_block_17->getOutput(0), calculate_padded_width(256, width_factor), 3, 2, 1, "model.18");
nvinfer1::ITensor* concat_inputs_19[] = { conv_block_18->getOutput(0), conv_block_14->getOutput(0) };
auto concat_layer_19 = model_network->addConcatenation(concat_inputs_19, 2);
auto c3_block_20 = C3(model_network, weights_map, *concat_layer_19->getOutput(0), calculate_padded_width(512, width_factor), calculate_padded_width(512, width_factor), calculate_padded_depth(3, depth_factor), false, 1, 0.5, "model.20");
// Camada YOLO 1
nvinfer1::IConvolutionLayer* detection_layer_1 = model_network->addConvolutionNd(*c3_block_20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), nvinfer1::DimsHW{ 1, 1 }, weights_map["model.24.m.1.weight"], weights_map["model.24.m.1.bias"]);
auto conv_block_21 = convBlock(model_network, weights_map, *c3_block_20->getOutput(0), calculate_padded_width(512, width_factor), 3, 2, 1, "model.21");
nvinfer1::ITensor* concat_inputs_22[] = { conv_block_21->getOutput(0), conv_block_10->getOutput(0) };
auto concat_layer_22 = model_network->addConcatenation(concat_inputs_22, 2);
auto c3_block_23 = C3(model_network, weights_map, *concat_layer_22->getOutput(0), calculate_padded_width(1024, width_factor), calculate_padded_width(1024, width_factor), calculate_padded_depth(3, depth_factor), false, 1, 0.5, "model.23");
nvinfer1::IConvolutionLayer* detection_layer_2 = model_network->addConvolutionNd(*c3_block_23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), nvinfer1::DimsHW{ 1, 1 }, weights_map["model.24.m.2.weight"], weights_map["model.24.m.2.bias"]);
auto yolo_output_layer = addYoLoLayer(model_network, weights_map, "model.24", std::vector<:iconvolutionlayer>{detection_layer_0, detection_layer_1, detection_layer_2});
yolo_output_layer->getOutput(0)->setName(OUTPUT_TENSOR_NAME);
model_network->markOutput(*yolo_output_layer->getOutput(0));
// Configurações e construção do motor
trt_builder->setMaxBatchSize(batch_size_limit);
builder_config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB de workspace
#if defined(ENABLE_FP16_MODE)
builder_config->setFlag(nvinfer1::BuilderFlag::kFP16);
#elif defined(ENABLE_INT8_MODE)
std::cout << "Sua plataforma suporta INT8: " << (trt_builder->platformHasFastInt8() ? "true" : "false") << std::endl;
assert(trt_builder->platformHasFastInt8());
builder_config->setFlag(nvinfer1::BuilderFlag::kINT8);
Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, MODEL_INPUT_WIDTH, MODEL_INPUT_HEIGHT, "./coco_calib/", "int8calib.table", INPUT_TENSOR_NAME);
builder_config->setInt8Calibrator(calibrator);
#endif
std::cout << "Construindo o motor do TensorRT, por favor aguarde..." << std::endl;
nvinfer1::ICudaEngine* cuda_engine = trt_builder->buildEngineWithConfig(*model_network, *builder_config);
std::cout << "Motor construído com sucesso!" << std::endl;
// Liberar recursos da rede
model_network->destroy();
// Liberar memória dos pesos no host
for (auto& entry : weights_map) {
free((void*)(entry.second.values));
}
return cuda_engine;
}
// Constrói o motor do TensorRT para o modelo YOLOv5-P6
nvinfer1::ICudaEngine* build_yolov5_p6_engine(
unsigned int batch_size_limit,
nvinfer1::IBuilder* trt_builder,
nvinfer1::IBuilderConfig* builder_config,
nvinfer1::DataType data_type,
float& depth_factor,
float& width_factor,
std::string& weights_file_path)
{
nvinfer1::INetworkDefinition* model_network = trt_builder->createNetworkV2(0U);
assert(model_network);
nvinfer1::ITensor* input_tensor = model_network->addInput(INPUT_TENSOR_NAME, data_type, nvinfer1::Dims3{ 3, MODEL_INPUT_HEIGHT, MODEL_INPUT_WIDTH });
assert(input_tensor);
std::map<:string weights=""> weights_map = loadWeights(weights_file_path);
/* ------ Backbone do YOLOv5-P6 ------ */
auto focus_layer_0 = focus(model_network, weights_map, *input_tensor, 3, calculate_padded_width(64, width_factor), 3, "model.0");
auto conv_block_1 = convBlock(model_network, weights_map, *focus_layer_0->getOutput(0), calculate_padded_width(128, width_factor), 3, 2, 1, "model.1");
auto c3_block_2 = C3(model_network, weights_map, *conv_block_1->getOutput(0), calculate_padded_width(128, width_factor), calculate_padded_width(128, width_factor), calculate_padded_depth(3, depth_factor), true, 1, 0.5, "model.2");
auto conv_block_3 = convBlock(model_network, weights_map, *c3_block_2->getOutput(0), calculate_padded_width(256, width_factor), 3, 2, 1, "model.3");
auto c3_block_4 = C3(model_network, weights_map, *conv_block_3->getOutput(0), calculate_padded_width(256, width_factor), calculate_padded_width(256, width_factor), calculate_padded_depth(9, depth_factor), true, 1, 0.5, "model.4");
auto conv_block_5 = convBlock(model_network, weights_map, *c3_block_4->getOutput(0), calculate_padded_width(512, width_factor), 3, 2, 1, "model.5");
auto c3_block_6 = C3(model_network, weights_map, *conv_block_5->getOutput(0), calculate_padded_width(512, width_factor), calculate_padded_width(512, width_factor), calculate_padded_depth(9, depth_factor), true, 1, 0.5, "model.6");
auto conv_block_7 = convBlock(model_network, weights_map, *c3_block_6->getOutput(0), calculate_padded_width(768, width_factor), 3, 2, 1, "model.7");
auto c3_block_8 = C3(model_network, weights_map, *conv_block_7->getOutput(0), calculate_padded_width(768, width_factor), calculate_padded_width(768, width_factor), calculate_padded_depth(3, depth_factor), true, 1, 0.5, "model.8");
auto conv_block_9 = convBlock(model_network, weights_map, *c3_block_8->getOutput(0), calculate_padded_width(1024, width_factor), 3, 2, 1, "model.9");
auto spp_block_10 = SPP(model_network, weights_map, *conv_block_9->getOutput(0), calculate_padded_width(1024, width_factor), calculate_padded_width(1024, width_factor), 3, 5, 7, "model.10");
auto c3_block_11 = C3(model_network, weights_map, *spp_block_10->getOutput(0), calculate_padded_width(1024, width_factor), calculate_padded_width(1024, width_factor), calculate_padded_depth(3, depth_factor), false, 1, 0.5, "model.11");
/* ------ Head do YOLOv5-P6 ------ */
auto conv_block_12 = convBlock(model_network, weights_map, *c3_block_11->getOutput(0), calculate_padded_width(768, width_factor), 1, 1, 1, "model.12");
auto upsample_layer_13 = model_network->addResize(*conv_block_12->getOutput(0));
assert(upsample_layer_13);
upsample_layer_13->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
upsample_layer_13->setOutputDimensions(c3_block_8->getOutput(0)->getDimensions());
nvinfer1::ITensor* concat_inputs_14[] = { upsample_layer_13->getOutput(0), c3_block_8->getOutput(0) };
auto concat_layer_14 = model_network->addConcatenation(concat_inputs_14, 2);
auto c3_block_15 = C3(model_network, weights_map, *concat_layer_14->getOutput(0), calculate_padded_width(1536, width_factor), calculate_padded_width(768, width_factor), calculate_padded_depth(3, depth_factor), false, 1, 0.5, "model.15");
auto conv_block_16 = convBlock(model_network, weights_map, *c3_block_15->getOutput(0), calculate_padded_width(512, width_factor), 1, 1, 1, "model.16");
auto upsample_layer_17 = model_network->addResize(*conv_block_16->getOutput(0));
assert(upsample_layer_17);
upsample_layer_17->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
upsample_layer_17->setOutputDimensions(c3_block_6->getOutput(0)->getDimensions());
nvinfer1::ITensor* concat_inputs_18[] = { upsample_layer_17->getOutput(0), c3_block_6->getOutput(0) };
auto concat_layer_18 = model_network->addConcatenation(concat_inputs_18, 2);
auto c3_block_19 = C3(model_network, weights_map, *concat_layer_18->getOutput(0), calculate_padded_width(1024, width_factor), calculate_padded_width(512, width_factor), calculate_padded_depth(3, depth_factor), false, 1, 0.5, "model.19");
auto conv_block_20 = convBlock(model_network, weights_map, *c3_block_19->getOutput(0), calculate_padded_width(256, width_factor), 1, 1, 1, "model.20");
auto upsample_layer_21 = model_network->addResize(*conv_block_20->getOutput(0));
assert(upsample_layer_21);
upsample_layer_21->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
upsample_layer_21->setOutputDimensions(c3_block_4->getOutput(0)->getDimensions());
nvinfer1::ITensor* concat_inputs_21[] = { upsample_layer_21->getOutput(0), c3_block_4->getOutput(0) };
auto concat_layer_22 = model_network->addConcatenation(concat_inputs_21, 2);
auto c3_block_23 = C3(model_network, weights_map, *concat_layer_22->getOutput(0), calculate_padded_width(512, width_factor), calculate_padded_width(256, width_factor), calculate_padded_depth(3, depth_factor), false, 1, 0.5, "model.23");
auto conv_block_24 = convBlock(model_network, weights_map, *c3_block_23->getOutput(0), calculate_padded_width(256, width_factor), 3, 2, 1, "model.24");
nvinfer1::ITensor* concat_inputs_25[] = { conv_block_24->getOutput(0), conv_block_20->getOutput(0) };
auto concat_layer_25 = model_network->addConcatenation(concat_inputs_25, 2);
auto c3_block_26 = C3(model_network, weights_map, *concat_layer_25->getOutput(0), calculate_padded_width(1024, width_factor), calculate_padded_width(512, width_factor), calculate_padded_depth(3, depth_factor), false, 1, 0.5, "model.26");
auto conv_block_27 = convBlock(model_network, weights_map, *c3_block_26->getOutput(0), calculate_padded_width(512, width_factor), 3, 2, 1, "model.27");
nvinfer1::ITensor* concat_inputs_28[] = { conv_block_27->getOutput(0), conv_block_16->getOutput(0) };
auto concat_layer_28 = model_network->addConcatenation(concat_inputs_28, 2);
auto c3_block_29 = C3(model_network, weights_map, *concat_layer_28->getOutput(0), calculate_padded_width(1536, width_factor), calculate_padded_width(768, width_factor), calculate_padded_depth(3, depth_factor), false, 1, 0.5, "model.29");
auto conv_block_30 = convBlock(model_network, weights_map, *c3_block_29->getOutput(0), calculate_padded_width(768, width_factor), 3, 2, 1, "model.30");
nvinfer1::ITensor* concat_inputs_31[] = { conv_block_30->getOutput(0), conv_block_12->getOutput(0) };
auto concat_layer_31 = model_network->addConcatenation(concat_inputs_31, 2);
auto c3_block_32 = C3(model_network, weights_map, *concat_layer_31->getOutput(0), calculate_padded_width(2048, width_factor), calculate_padded_width(1024, width_factor), calculate_padded_depth(3, depth_factor), false, 1, 0.5, "model.32");
/* ------ Camadas de Detecção ------ */
nvinfer1::IConvolutionLayer* detection_layer_0_p6 = model_network->addConvolutionNd(*c3_block_23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), nvinfer1::DimsHW{ 1, 1 }, weights_map["model.33.m.0.weight"], weights_map["model.33.m.0.bias"]);
nvinfer1::IConvolutionLayer* detection_layer_1_p6 = model_network->addConvolutionNd(*c3_block_26->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), nvinfer1::DimsHW{ 1, 1 }, weights_map["model.33.m.1.weight"], weights_map["model.33.m.1.bias"]);
nvinfer1::IConvolutionLayer* detection_layer_2_p6 = model_network->addConvolutionNd(*c3_block_29->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), nvinfer1::DimsHW{ 1, 1 }, weights_map["model.33.m.2.weight"], weights_map["model.33.m.2.bias"]);
nvinfer1::IConvolutionLayer* detection_layer_3_p6 = model_network->addConvolutionNd(*c3_block_32->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), nvinfer1::DimsHW{ 1, 1 }, weights_map["model.33.m.3.weight"], weights_map["model.33.m.3.bias"]);
auto yolo_output_layer_p6 = addYoLoLayer(model_network, weights_map, "model.33", std::vector<:iconvolutionlayer>{detection_layer_0_p6, detection_layer_1_p6, detection_layer_2_p6, detection_layer_3_p6});
yolo_output_layer_p6->getOutput(0)->setName(OUTPUT_TENSOR_NAME);
model_network->markOutput(*yolo_output_layer_p6->getOutput(0));
// Configurações e construção do motor
trt_builder->setMaxBatchSize(batch_size_limit);
builder_config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB de workspace
#if defined(ENABLE_FP16_MODE)
builder_config->setFlag(nvinfer1::BuilderFlag::kFP16);
#elif defined(ENABLE_INT8_MODE)
std::cout << "Sua plataforma suporta INT8: " << (trt_builder->platformHasFastInt8() ? "true" : "false") << std::endl;
assert(trt_builder->platformHasFastInt8());
builder_config->setFlag(nvinfer1::BuilderFlag::kINT8);
Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, MODEL_INPUT_WIDTH, MODEL_INPUT_HEIGHT, "./coco_calib/", "int8calib.table", INPUT_TENSOR_NAME);
builder_config->setInt8Calibrator(calibrator);
#endif
std::cout << "Construindo o motor do TensorRT (P6), por favor aguarde..." << std::endl;
nvinfer1::ICudaEngine* cuda_engine = trt_builder->buildEngineWithConfig(*model_network, *builder_config);
std::cout << "Motor P6 construído com sucesso!" << std::endl;
// Liberar recursos da rede
model_network->destroy();
// Liberar memória dos pesos no host
for (auto& entry : weights_map) {
free((void*)(entry.second.values));
}
return cuda_engine;
}
// Cria e serializa o motor do TensorRT para uso posterior
void create_and_serialize_engine(unsigned int batch_size_limit, nvinfer1::IHostMemory** serialized_engine_stream, float& depth_factor, float& width_factor, std::string& weights_path) {
// Cria o builder do TensorRT
nvinfer1::IBuilder* trt_builder = nvinfer1::createInferBuilder(trt_logger);
nvinfer1::IBuilderConfig* builder_config = trt_builder->createBuilderConfig();
// Constrói o motor de inferência
nvinfer1::ICudaEngine* cuda_engine = build_yolov5_engine(batch_size_limit, trt_builder, builder_config, nvinfer1::DataType::kFLOAT, depth_factor, width_factor, weights_path);
assert(cuda_engine != nullptr);
// Serializa o motor para um buffer
(*serialized_engine_stream) = cuda_engine->serialize();
// Limpar objetos
cuda_engine->destroy();
trt_builder->destroy();
builder_config->destroy();
}
// Realiza a inferência assíncrona na GPU
void execute_inference(nvinfer1::IExecutionContext& inference_context, cudaStream_t& cuda_stream, void** device_buffers, float* host_input_data, float* host_output_data, int current_batch_size) {
// Copia os dados de entrada do host para o dispositivo (GPU)
CUDA_CHECK(cudaMemcpyAsync(device_buffers[0], host_input_data, current_batch_size * 3 * MODEL_INPUT_HEIGHT * MODEL_INPUT_WIDTH * sizeof(float), cudaMemcpyHostToDevice, cuda_stream));
// Enfileira a tarefa de inferência na GPU
inference_context.enqueue(current_batch_size, device_buffers, cuda_stream, nullptr);
// Copia os resultados de saída do dispositivo para o host
CUDA_CHECK(cudaMemcpyAsync(host_output_data, device_buffers[1], current_batch_size * DETECTION_OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, cuda_stream));
// Sincroniza o stream CUDA para garantir que a inferência e as cópias estejam completas
cudaStreamSynchronize(cuda_stream);
}
// Analisa os argumentos da linha de comando
bool parse_command_line_args(int argument_count, char** argument_values, std::string& engine_file_path) {
if (argument_count < 3) return false;
if (std::string(argument_values[1]) == "-v" && argument_count == 3) {
engine_file_path = std::string(argument_values[2]);
} else {
return false;
}
return true;
}
int main(int argument_count, char** argument_values) {
cudaSetDevice(GPU_DEVICE_ID);
std::string engine_file_path = "";
if (!parse_command_line_args(argument_count, argument_values, engine_file_path)) {
std::cerr << "Argumentos inválidos!" << std::endl;
std::cerr << "Uso: ./yolov5 -v [caminho_para_engine.engine] // Executa inferência com câmera" << std::endl;
return -1;
}
std::ifstream engine_file_handle(engine_file_path, std::ios::binary);
if (!engine_file_handle.good()) {
std::cerr << "Erro ao ler o arquivo: " << engine_file_path << std::endl;
return -1;
}
char* serialized_model_data{ nullptr };
size_t data_size = 0;
engine_file_handle.seekg(0, engine_file_handle.end);
data_size = engine_file_handle.tellg();
engine_file_handle.seekg(0, engine_file_handle.beg);
serialized_model_data = new char[data_size];
assert(serialized_model_data);
engine_file_handle.read(serialized_model_data, data_size);
engine_file_handle.close();
// Preparar buffers de entrada/saída
static float input_data_buffer[INFERENCE_BATCH_SIZE * 3 * MODEL_INPUT_HEIGHT * MODEL_INPUT_WIDTH];
static float output_probabilities[INFERENCE_BATCH_SIZE * DETECTION_OUTPUT_SIZE];
nvinfer1::IRuntime* inference_runtime = nvinfer1::createInferRuntime(trt_logger);
assert(inference_runtime != nullptr);
nvinfer1::ICudaEngine* loaded_engine = inference_runtime->deserializeCudaEngine(serialized_model_data, data_size);
assert(loaded_engine != nullptr);
nvinfer1::IExecutionContext* execution_context = loaded_engine->createExecutionContext();
assert(execution_context != nullptr);
delete[] serialized_model_data;
assert(loaded_engine->getNbBindings() == 2);
void* device_buffers[2];
const int input_binding_idx = loaded_engine->getBindingIndex(INPUT_TENSOR_NAME);
const int output_binding_idx = loaded_engine->getBindingIndex(OUTPUT_TENSOR_NAME);
assert(input_binding_idx == 0);
assert(output_binding_idx == 1);
// Alocar buffers na GPU
CUDA_CHECK(cudaMalloc(&device_buffers[input_binding_idx], INFERENCE_BATCH_SIZE * 3 * MODEL_INPUT_HEIGHT * MODEL_INPUT_WIDTH * sizeof(float)));
CUDA_CHECK(cudaMalloc(&device_buffers[output_binding_idx], INFERENCE_BATCH_SIZE * DETECTION_OUTPUT_SIZE * sizeof(float)));
// Criar stream CUDA
cudaStream_t cuda_stream;
CUDA_CHECK(cudaStreamCreate(&cuda_stream));
// Inicializar captura de vídeo
cv::VideoCapture video_capture(0); // Tente 0 para webcam padrão, se falhar, tente 1 ou outro ID
if (!video_capture.isOpened()) {
std::cerr << "Erro ao abrir o fluxo de vídeo da câmera." << std::endl;
return -1;
}
int keyboard_input_key;
int frame_processing_count = 0;
cv::Mat current_frame;
while (true) {
video_capture >> current_frame;
if (current_frame.empty()) {
std::cerr << "Falha ao ler o frame da câmera!" << std::endl;
break;
}
frame_processing_count++;
// Pré-processamento do frame
for (int b = 0; b < frame_processing_count; b++) {
cv::Mat preprocessed_image = preprocess_img(current_frame, MODEL_INPUT_WIDTH, MODEL_INPUT_HEIGHT); // letterbox BGR para RGB
int buffer_index = 0;
for (int row_idx = 0; row_idx < MODEL_INPUT_HEIGHT; ++row_idx) {
uchar* pixel_ptr = preprocessed_image.data + row_idx * preprocessed_image.step;
for (int col_idx = 0; col_idx < MODEL_INPUT_WIDTH; ++col_idx) {
input_data_buffer[b * 3 * MODEL_INPUT_HEIGHT * MODEL_INPUT_WIDTH + buffer_index] = (float)pixel_ptr[2] / 255.0; // R
input_data_buffer[b * 3 * MODEL_INPUT_HEIGHT * MODEL_INPUT_WIDTH + buffer_index + MODEL_INPUT_HEIGHT * MODEL_INPUT_WIDTH] = (float)pixel_ptr[1] / 255.0; // G
input_data_buffer[b * 3 * MODEL_INPUT_HEIGHT * MODEL_INPUT_WIDTH + buffer_index + 2 * MODEL_INPUT_HEIGHT * MODEL_INPUT_WIDTH] = (float)pixel_ptr[0] / 255.0; // B
pixel_ptr += 3;
++buffer_index;
}
}
}
// Executar inferência
auto start_time = std::chrono::system_clock::now();
execute_inference(*execution_context, cuda_stream, device_buffers, input_data_buffer, output_probabilities, INFERENCE_BATCH_SIZE);
auto end_time = std::chrono::system_clock::now();
// Calcular FPS
long long elapsed_ms = std::chrono::duration_cast<:chrono::milliseconds>(end_time - start_time).count();
double frames_per_second = 1000.0 / elapsed_ms;
// Processar resultados da detecção
std::vector<:vector>> detection_results_batch(frame_processing_count);
for (int b = 0; b < frame_processing_count; b++) {
auto& current_frame_detections = detection_results_batch[b];
nms(current_frame_detections, &output_probabilities[b * DETECTION_OUTPUT_SIZE], CONFIDENCE_THRESHOLD_VAL, NMS_THRESHOLD_VAL);
}
// Desenhar caixas delimitadoras e etiquetas no frame
for (int b = 0; b < frame_processing_count; b++) {
auto& detections_to_draw = detection_results_batch[b];
for (size_t j = 0; j < detections_to_draw.size(); j++) {
cv::Rect bounding_box_rect = get_rect(current_frame, detections_to_draw[j].bbox);
cv::rectangle(current_frame, bounding_box_rect, cv::Scalar(0x27, 0xC1, 0x36), 2);
std::string class_label_text = object_classes[(int)detections_to_draw[j].class_id];
cv::putText(current_frame, class_label_text, cv::Point(bounding_box_rect.x, bounding_box_rect.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
}
std::string fps_display_text = "FPS: " + std::to_string((int)frames_per_second);
cv::putText(current_frame, fps_display_text, cv::Point(11, 80), cv::FONT_HERSHEY_PLAIN, 3, cv::Scalar(0, 0, 255), 2, cv::LINE_AA);
}
cv::imshow("Inferência YOLOv5 com TensorRT", current_frame);
keyboard_input_key = cv::waitKey(1);
if (keyboard_input_key == 'q') {
break;
}
frame_processing_count = 0; // Reset para processar um frame por vez
}
// Liberar recursos
video_capture.release();
cudaStreamDestroy(cuda_stream);
CUDA_CHECK(cudaFree(device_buffers[input_binding_idx]));
CUDA_CHECK(cudaFree(device_buffers[output_binding_idx]));
execution_context->destroy();
loaded_engine->destroy();
inference_runtime->destroy();
return 0;
}
</:vector></:chrono::milliseconds></:iconvolutionlayer></:string></:iconvolutionlayer></:string>
Após substituir o código, recompile o projeto:
make
Finalmente, execute o programa para iniciar a detecção de objetos via câmera:
./yolov5 -v yolov5s.engine
Isso concluirá a configuração e o teste do TensorRT com YOLOv5, permitindo a inferência acelerada em tempo real.