semantic_lstm_vehicle_tensorrt.h

浏览该文件的文档.

/******************************************************************************
 * Copyright 2023 The Apollo Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *****************************************************************************/
 
#pragma once
 
#include <string>
#include <vector>
#include <memory>
#include <mutex>
#include <map>
#include <utility>
 
#include "NvInfer.h"
#include "NvInferRuntime.h"
 
#include "cyber/plugin_manager/plugin_manager.h"
#include "modules/prediction/evaluator/model_manager/model/model_base.h"
 
namespace apollo {
namespace prediction {
 
class SemanticLstmVehicleTensorrt : public ModelBase {
public:
    SemanticLstmVehicleTensorrt() {}
    ~SemanticLstmVehicleTensorrt() {
        Destory();
    }
 
    virtual bool Init();
 
    virtual bool Inference(
            const std::vector<void*>& input_buffer,
            unsigned int input_size,
            std::vector<void*>* output_buffer,
            unsigned int output_size);
 
    virtual bool LoadModel();
 
    virtual void Destory();
 
private:
    nvinfer1::ICudaEngine* engine_ = nullptr;
    nvinfer1::IRuntime* runtime_ = nullptr;
    nvinfer1::IExecutionContext* context_ = nullptr;
    std::vector<void*> buffers_;
    cudaStream_t stream_;
 
    // input output name based on onnx model
    const char* INPUT_IMAGE_NAME = "img_tensor";
    const char* INPUT_OBSTACLE_POS_NAME = "obstacle_pos";
    const char* INPUT_OBSTACLE_POS_STEP_NAME = "obstacle_pos_step";
    const char* OUTPUT_NAME = "predict";
    // input output index
    int input_image_index_;
    int input_obstacle_pos_index_;
    int input_obstacle_pos_step_index_;
    int output_index_;
    std::mutex mtx;
 
    class RTLogger : public nvinfer1::ILogger {
        void log(Severity severity, const char* msg) noexcept override {
            switch (severity) {
            case Severity::kINTERNAL_ERROR:
            case Severity::kERROR:
                AERROR << msg;
                break;
            case Severity::kWARNING:
                AWARN << msg;
                break;
            case Severity::kINFO:
            case Severity::kVERBOSE:
                ADEBUG << msg;
                break;
            default:
                break;
            }
        }
    } rt_gLogger;
};
CYBER_PLUGIN_MANAGER_REGISTER_PLUGIN(apollo::prediction::SemanticLstmVehicleTensorrt, ModelBase)
 
}  // namespace prediction
}  // namespace apollo