TensorRT 학습(2): Pytorch 모델에서 TensorRT C++ 프로세스로

머리말

  https://github.com/wang-xinyu/pytorchx
  https://github.com/wang-xinyu/tensorrtx
  이 기사는 TensorRT의 C++ API를 배우기 위해 lenet과 mlp의 두 가지 간단한 예제를 선택합니다.
  코드 부분은 assert판단 및 destroy해제 공간 부분을 건너뛰고 더 많은 핵심 부분만 유지합니다.

1. pytorxx

model.py와   pytorchx 에는 두 개의 파일이 있습니다 inference.py. model.pypytorch 모델을 구축하고 가중치를 저장하는 데 사용 model.pth되며 .   이것은 복잡하지 않으며 C++에서 모델의 가중치를 읽기 위해 텍스트 형식으로 모델의 가중치를 저장합니다.inference.pymodel.pthmodel.wts
model.wts

# mlp
2
linear.weight 1 3fffdecf
linear.bias 1 3b16e580

# lenet
10
conv1.weight 150 be40ee1b bd20baba ...
conv1.bias 6 bd32705a 3e2182a8 ...
conv2.weight 2400 3c6f2224 3c69308f ...
conv2.bias 16 bd183967 bcb1ac89 ...
fc1.weight 48000 3c162c20 bd25196a ...
fc1.bias 120 3d3c3d4a bc64b947 ...
fc2.weight 10080 bce095a3 3d33b9dc ...
fc2.bias 84 bc71eaa0 3d9b276d ...
fc3.weight 840 3c25286d 3d855351 ...
fc3.bias 10 bdbe4bb8 3b119ed1 ...

  코드가 더 쉽고 분석되지 않으며 고통스러운 C++ 링크가 이어집니다.

2. 텐소르텍스

2.1 cmake

  당분간 이 과정에 대한 심도 있는 연구는 하지 않겠습니다.핵심은 CMakeLists.txt내용이 비슷하고 메인 내용은 화려한 model.cpp내용입니다. ./model -s엔진 생성 model.engine및 ./model -d추론 실행 의 두 부분을 보려면 여기를 클릭하십시오 .

2.2 생성 엔진

2.2.1 주요 프로세스

  이 두 모델의 주요 프로세스는 동일합니다. Create modelStream $\to$APIToModel 모델 구축→$\to$ 쓰다model.engine

IHostMemory* modelStream{
    
    nullptr};
APIToModel(1, &modelStream);
std::ofstream p("../model.engine", std::ios::binary);
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());

2.2.2APIToModel

  프로세스는 여전히 동일하며 구성 정보를 logging.h읽고 모델을 생성합니다.이 헤더 파일에 대한 작성자의 설명은 NVIDIA TRT API를 사용하기 위한 로거 파일(대부분 모든 모델에 동일)이며 직접 사용해야 합니다. 건너뛰기 지금은. 에 핵심이 있음을 알 수 있습니다 createModelEngine.

void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream){
    
    
	// Create builder with the help of logger
	IBuilder *builder = createInferBuilder(gLogger);
	// Create hardware configs
	IBuilderConfig *config = builder->createBuilderConfig();
	// Build an engine
	ICudaEngine* engine = createModelEngine(maxBatchSize, builder, config, DataType::kFLOAT);
	assert(engine != nullptr);
	// serialize the engine into binary stream
	(*modelStream) = engine->serialize();
	// free up the memory
    engine->destroy();
    builder->destroy();
}

2.2.3createModelEngine

(1)
  프로세스의 네트워크 부분을 구축하는 것 외에 작업 공간 크기의 설정 만 다릅니다. 공식 블로그에 설명이 있습니다 setMaxWorkspaceSize. 대략적인 GPU 공간 할당은 1ULL << 301GB이며 단위는 할당 공간은 바이트입니다.

"两个函数的输入是相同的"
ICudaEngine* createMLPEngine(unsigned int maxBatchSize, IBuilder *builder, IBuilderConfig *config, DataType dt)
ICudaEngine* createLenetEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt)

"读取权重, wts文件的格式是相同的, 估计是通用的, 读进来是个字典, 跳过读取细节"
std::map<std::string, Weights> weightMap = loadWeights("../mlp.wts");
std::map<std::string, Weights> weightMap = loadWeights("../lenet5.wts");
"创建空模型"
INetworkDefinition* network = builder->createNetworkV2(0U);
INetworkDefinition* network = builder->createNetworkV2(0U);
"创建输入, name 和 type 是一样的, lenet的维度是 1, 28, 28"
"这里看起来比pytorch的输入少了一个batchsize的维度, 但在后面有 setMaxBatchSize"
ITensor* data = network->addInput("data", DataType::kFLOAT, Dims3{
    
    1, 1, 1});
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{
    
    1, INPUT_H, INPUT_W});
"构建网络, 后面单独说明"

"Set configurations"
builder->setMaxBatchSize(1);
builder->setMaxBatchSize(maxBatchSize);
"Set workspace size"
config->setMaxWorkspaceSize(1 << 20);
config->setMaxWorkspaceSize(16 << 20);
"Build CUDA Engine using network and configurations"
ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);

(2) 네트워크 부분 구축

// mlp
IFullyConnectedLayer *fc1 = network->addFullyConnected(*data, 1, weightMap["linear.weight"], weightMap["linear.bias"]);
fc1->getOutput(0)->setName("out");
network->markOutput(*fc1->getOutput(0));

  mlp의 구조는 너무 간단합니다. lenet만 보세요. 먼저 lenet의 구조를 검토합니다:
  $[1,1,32,32]$ $\to$ Conv2d(1,6,5,1,0)+relu $\to$ $[1,6,28,\_28]$ $\to$ AvgPool2d(2,2,0) $\to$ $[1,6,14,\_14]$ $\to$ Conv2d(6,16,5,1,0)+relu $\to$ $[1,16,\_10\_\_10]$ $\to$ AvgPool2d(2,2,0) $\to$ $[1,16,\_5,5]$ $\to$ $[1,400]$ $\to$ Linear(400,120)+relu $\to$ $[1,120]$ $\to$ Linear(120,84)+relu $\to$ $[1,84]$ $\to$ Linear(84,10)+softmax $\to$ $[1,10]\_$

  pytorch의 API에 해당하는 tensorrt API의 사용법을 볼 수 있습니다.
Convolution: addConvolutionNd(输入, 输出维度, kernel大小, 权重, 偏置)
Activation: addActivation(输入, 激活函数类型)
Pooling: addPoolingNd(输入, 池化模式, kernel大小)
setStrideNdconvolution 및 pooling의 단계 크기를 설정하는 데 사용됩니다. 전체
연결:addFullyConnected(输入, 输出维度, 权重, 偏置)

  마지막으로 마지막 레이어의 출력에 이름을 지정하고 네트워크가 이를 출력으로 표시하도록 합니다.

// lenet
IConvolutionLayer* conv1 = network->addConvolutionNd(*data, 6, DimsHW{
    
    5, 5}, weightMap["conv1.weight"], weightMap["conv1.bias"]);
conv1->setStrideNd(DimsHW{
    
    1, 1});
IActivationLayer* relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU);
IPoolingLayer* pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kAVERAGE, DimsHW{
    
    2, 2});
pool1->setStrideNd(DimsHW{
    
    2, 2});

IConvolutionLayer* conv2 = network->addConvolutionNd(*pool1->getOutput(0), 16, DimsHW{
    
    5, 5}, weightMap["conv2.weight"], weightMap["conv2.bias"]);
conv2->setStrideNd(DimsHW{
    
    1, 1});
IActivationLayer* relu2 = network->addActivation(*conv2->getOutput(0), ActivationType::kRELU);
IPoolingLayer* pool2 = network->addPoolingNd(*relu2->getOutput(0), PoolingType::kAVERAGE, DimsHW{
    
    2, 2});
pool2->setStrideNd(DimsHW{
    
    2, 2});

IFullyConnectedLayer* fc1 = network->addFullyConnected(*pool2->getOutput(0), 120, weightMap["fc1.weight"], weightMap["fc1.bias"]);
IActivationLayer* relu3 = network->addActivation(*fc1->getOutput(0), ActivationType::kRELU);
IFullyConnectedLayer* fc2 = network->addFullyConnected(*relu3->getOutput(0), 84, weightMap["fc2.weight"], weightMap["fc2.bias"]);
IActivationLayer* relu4 = network->addActivation(*fc2->getOutput(0), ActivationType::kRELU);
IFullyConnectedLayer* fc3 = network->addFullyConnected(*relu4->getOutput(0), OUTPUT_SIZE, weightMap["fc3.weight"], weightMap["fc3.bias"]);

ISoftMaxLayer* prob = network->addSoftMax(*fc3->getOutput(0));
prob->getOutput(0)->setName(OUTPUT_BLOB_NAME);
network->markOutput(*prob->getOutput(0));

2.3 추론 실행

2.3.1 주요 프로세스

  를 제외한 전체 추론 과정은 doInferencetensorrt API를 사용하며, 기성품 logging.h과 필요한 것만 알 수 있습니다 model.engine.

// create a model using the API directly and serialize it to a stream
char *trtModelStream{
    
    nullptr};
size_t size{
    
    0};
// read model from the engine file
std::ifstream file("../model.engine", std::ios::binary);
if (file.good()) {
    
    
    file.seekg(0, file.end);
    size = file.tellg();
    file.seekg(0, file.beg);
    trtModelStream = new char[size];
    assert(trtModelStream);
    file.read(trtModelStream, size);
    file.close();
}
// create a runtime (required for deserialization of model) with NVIDIA's logger
IRuntime *runtime = createInferRuntime(gLogger);
// deserialize engine for using the char-stream
                               deserializeCudaEngine
ICudaEngine *engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
// create execution context -- required for inference executions
IExecutionContext *context = engine->createExecutionContext();
// 创建输入输出
float data[INPUT_SIZE];		// mlp:1	lenet:H*W
float out[OUTPUT_SIZE];		// mlp:1	lenet:10

// time the execution
auto start = std::chrono::system_clock::now();
// do inference using the parameters
doInference(*context, data, out, 1);
// time the execution
auto end = std::chrono::system_clock::now();

2.3.2doInference

  lenet과 mlp doInference도 같고, 추론은 사실 한 문장에 불과하고 context.enqueue(batchSize, buffers, stream, nullptr);, 나머지는 입력과 출력 공간을 만들고 저장하고, cuda 스트림을 만들고, CPU와 GPU 간에 데이터를 이동하는 작업을 합니다.
  yolov5의 코드를 보면 Running Reasoning은 사실 한 문장에 불과하고 주된 작업은 데이터를 준비하는 것입니다.

void doInference(IExecutionContext &context, float *input, float *output, int batchSize) {
    
    
    // Get engine from the context
    const ICudaEngine &engine = context.getEngine();

    // Pointers to input and output device buffers to pass to engine.
    // Engine requires exactly IEngine::getNbBindings() number of buffers.
    assert(engine.getNbBindings() == 2);
    void *buffers[2];

    // In order to bind the buffers, we need to know the names of the input and output tensors.
    // Note that indices are guaranteed to be less than IEngine::getNbBindings()
    const int inputIndex = engine.getBindingIndex("data");
    const int outputIndex = engine.getBindingIndex("out");

    // Create GPU buffers on device -- allocate memory for input and output
    cudaMalloc(&buffers[inputIndex], batchSize * INPUT_SIZE * sizeof(float));
    cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float));

    // create CUDA stream for simultaneous CUDA operations
    cudaStream_t stream;
    cudaStreamCreate(&stream);

    // copy input from host (CPU) to device (GPU)  in stream
    cudaMemcpyAsync(buffers[inputIndex], input, batchSize * INPUT_SIZE * sizeof(float), cudaMemcpyHostToDevice, stream);

    // execute inference using context provided by engine
    context.enqueue(batchSize, buffers, stream, nullptr);

    // copy output back from device (GPU) to host (CPU)
    cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost,
                    stream);

    // synchronize the stream to prevent issues
    //      (block CUDA and wait for CUDA operations to be completed)
    cudaStreamSynchronize(stream);

    // Release stream and buffers (memory)
    cudaStreamDestroy(stream);
    cudaFree(buffers[inputIndex]);
    cudaFree(buffers[outputIndex]);
}