Tensorrt-mnist_caffe

准备数据：

在tensorrt安装后，从/usr/src/tensorrt/data下载pgms

export TRT_DATADIR=/usr/src/tensorrt/data
pushd $TRT_DATADIR/mnist
pip install Pillow
python3 download_pgms.py
popd

运行sample：

#include "common/argsParser.h"
#include "common/buffers.h"
#include "common/common.h"
#include "common/logger.h"

#include "NvCaffeParser.h"
#include "NvInfer.h"

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cuda_runtime_api.h>
#include <fstream>
#include <iostream>
#include <sstream>
const std::string gSampleName = "TensorRT.sample_mnist";
class SampleMNIST
{
    template<typename T>
    using  SampleUniquePtr = std::unique_ptr<T,samplesCommon::InferDeleter>;
public:
    SampleMNIST(const samplesCommon::CaffeSampleParams& params)
        :mParams(params)
    {

    }
    bool build();
    bool infer();
    bool teardown();

private:
    bool constructNetwork(SampleUniquePtr<nvcaffeparser1::ICaffeParser>& parser,SampleUniquePtr<nvinfer1::INetworkDefinition>& network);
    bool processInput(const samplesCommon::BufferManager& buffers,const std::string& inputTensorName,int inputFileIndex) const;
    bool verifyOutput(const samplesCommon::BufferManager& buffers,const std::string& outputTensorName, int groundTruthDigit)const;
    std::shared_ptr<nvinfer1::ICudaEngine> mEngine{nullptr};
    samplesCommon::CaffeSampleParams mParams;
    nvinfer1::Dims mInputDims;
    SampleUniquePtr<nvcaffeparser1::IBinaryProtoBlob> mMeanBlob;

};
bool SampleMNIST::build()
{
    auto builder = SampleUniquePtr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
    if(!builder)
    {
        return false;
    }
    auto network = SampleUniquePtr<nvinfer1::INetworkDefinition>(builder->createNetwork());
    if(!network)
    {
        return false;
    }
    auto config = SampleUniquePtr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
    if(!config)
    {
        return false;
    }
    auto parser = SampleUniquePtr<nvcaffeparser1::ICaffeParser>(nvcaffeparser1::createCaffeParser());
    if(!parser)
    {
        return false;
    }
    if(!constructNetwork(parser,network))
    {
        return false;
    }
    builder->setMaxBatchSize(mParams.batchSize);
    config->setMaxWorkspaceSize(16_MiB);
    config->setFlag(BuilderFlag::kGPU_FALLBACK);
    config->setFlag(BuilderFlag::kSTRICT_TYPES);
    if(mParams.fp16)
    {
        config->setFlag(BuilderFlag::kFP16);
    }
    if(mParams.int8)
    {
        config->setFlag(BuilderFlag::kINT8);
    }
    samplesCommon::enableDLA(builder.get(),config.get(),mParams.dlaCore);

    mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(builder->buildEngineWithConfig(*network,*config),samplesCommon::InferDeleter());
    if(!mEngine)
    {
        return false;
    }
    assert(network->getNbInputs() == 1);
    mInputDims = network->getInput(0)->getDimensions();
    assert(mInputDims.nbDims == 3);

    return true;
}
bool SampleMNIST::constructNetwork(SampleUniquePtr<nvcaffeparser1::ICaffeParser>& parser, SampleUniquePtr<nvinfer1::INetworkDefinition>& network)
{
    const nvcaffeparser1::IBlobNameToTensor* blobNameToTensor = parser->parse(
                mParams.prototxtFileName.c_str(),mParams.weightsFileName.c_str(),*network,nvinfer1::DataType::kFLOAT);
    for(auto& s : mParams.outputTensorNames)
    {
        network->markOutput(*blobNameToTensor->find(s.c_str()));
    }

    nvinfer1::Dims inputDims = network->getInput(0)->getDimensions();
    mMeanBlob=
            SampleUniquePtr<nvcaffeparser1::IBinaryProtoBlob>(parser->parseBinaryProto(mParams.meanFileName.c_str()));
    nvinfer1::Weights meanWeights{nvinfer1::DataType::kFLOAT,mMeanBlob->getData(),inputDims.d[1]*inputDims.d[1]};
    float maxMean
        = samplesCommon::getMaxValue(static_cast<const float*>(meanWeights.values), samplesCommon::volume(inputDims));

    auto mean = network->addConstant(nvinfer1::Dims3(1, inputDims.d[1], inputDims.d[2]), meanWeights);
    if (!mean->getOutput(0)->setDynamicRange(-maxMean, maxMean))
    {
        return false;
    }
    if (!network->getInput(0)->setDynamicRange(-maxMean, maxMean))
    {
        return false;
    }
    auto meanSub = network->addElementWise(*network->getInput(0), *mean->getOutput(0), ElementWiseOperation::kSUB);
    if (!meanSub->getOutput(0)->setDynamicRange(-maxMean, maxMean))
    {
        return false;
    }
    network->getLayer(0)->setInput(0, *meanSub->getOutput(0));
    samplesCommon::setAllTensorScales(network.get(), 127.0f, 127.0f);

    return true;

}
bool SampleMNIST::infer()
{
    samplesCommon::BufferManager buffers(mEngine,mParams.batchSize);
    auto context = SampleUniquePtr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
    if(!context)
    {
        return false;
    }
    srand(time(NULL));
    const int digit = rand() % 10;
    assert(mParams.inputTensorNames.size() == 1);
    if (!processInput(buffers, mParams.inputTensorNames[0], digit))
    {
        return false;
    }
    cudaStream_t stream;
    CHECK(cudaStreamCreate(&stream));
    buffers.copyInputToDeviceAsync(stream);
    if(!context->enqueue(mParams.batchSize, buffers.getDeviceBindings().data(),stream,nullptr))
    {
        return false;
    }
    buffers.copyOutputToHostAsync(stream);
    cudaStreamSynchronize(stream);
    cudaStreamDestroy(stream);
    assert(mParams.outputTensorNames.size()==1);
    bool outputCorrect = verifyOutput(buffers, mParams.outputTensorNames[0], digit);

    return outputCorrect;
}
bool SampleMNIST::processInput(const samplesCommon::BufferManager& buffers, const std::string& inputTensorName, int inputFileIdx) const
{
    const int inputH = mInputDims.d[1];
    const int inputW = mInputDims.d[2];

    srand(unsigned(time(nullptr)));
    std::vector<uint8_t> fileData(inputH*inputW);
    readPGMFile(locateFile(std::to_string(inputFileIdx) + ".pgm",mParams.dataDirs),fileData.data(),inputH,inputW);

    sample::gLogInfo<<"Input:\n";
    for (int i = 0; i < inputH * inputW; i++)
    {
        sample::gLogInfo << (" .:-=+*#%@"[fileData[i] / 26]) << (((i + 1) % inputW) ? "" : "\n");
    }
    sample::gLogInfo << std::endl;

    float* hostInputBuffer = static_cast<float*>(buffers.getHostBuffer(inputTensorName));
    for(int i=0;i<inputH*inputW;i++)
    {
        hostInputBuffer[i]= float(fileData[i]);
    }
    return true;

}
bool SampleMNIST::verifyOutput(const samplesCommon::BufferManager& buffers, const std::string& outputTensorName, int groundTruthDigit) const
{
    const float* prob= static_cast<const float*>(buffers.getHostBuffer(outputTensorName));

    sample::gLogInfo <<"Output:\n";
    float val{0.0f};
    int idx{0};
    const int kDIGITS=10;
    for(int i=0;i<kDIGITS;i++)
    {
        if(val<prob[i])
        {
            val=prob[i];
            idx=i;
        }
        sample::gLogInfo<<i<<" : "<<std::string(int(std::floor(prob[i]*10 + 0.5f)),'*')<<"\n";
    }
    sample::gLogInfo <<std::endl;
    return (idx == groundTruthDigit && val > 0.9f);
}
bool SampleMNIST::teardown()
{
    nvcaffeparser1::shutdownProtobufLibrary();
    return true;
}
samplesCommon::CaffeSampleParams initializeSampleParams(const samplesCommon::Args& args)
{
    samplesCommon::CaffeSampleParams params;
    if(args.dataDirs.empty())
    {
        params.dataDirs.push_back("./data/");
//        params.dataDirs.push_back("./data/");
    }
    else
    {
        params.dataDirs = args.dataDirs;
    }

    params.prototxtFileName = locateFile("mnist.prototxt",params.dataDirs);
    params.weightsFileName = locateFile("mnist.caffemodel",params.dataDirs);
    params.meanFileName = locateFile("mnist_mean.binaryproto",params.dataDirs);
    params.inputTensorNames.push_back("data");
    params.batchSize=1;
    params.outputTensorNames.push_back("prob");
    params.dlaCore=args.useDLACore;
    params.int8=args.runInInt8;
    params.fp16=args.runInFp16;
    return params;

}
void printHelpInfo()
{
    std::cout
        << "Usage: ./sample_mnist [-h or --help] [-d or --datadir=<path to data directory>] [--useDLACore=<int>]\n";
    std::cout << "--help          Display help information\n";
    std::cout << "--datadir       Specify path to a data directory, overriding the default. This option can be used "
                 "multiple times to add multiple directories. If no data directories are given, the default is to use "
                 "(data/samples/mnist/, data/mnist/)"
              << std::endl;
    std::cout << "--useDLACore=N  Specify a DLA engine for layers that support DLA. Value can range from 0 to n-1, "
                 "where n is the number of DLA engines on the platform."
              << std::endl;
    std::cout << "--int8          Run in Int8 mode.\n";
    std::cout << "--fp16          Run in FP16 mode.\n";
}
int main(int argc,char** argv)
{
    samplesCommon::Args args;
    bool argsOK= samplesCommon::parseArgs(args,argc,argv);
    if(!argsOK)
    {
        sample::gLogError << "Invalid arguments" << std::endl;
        printHelpInfo();
        return EXIT_FAILURE;
    }
    if (args.help)
    {
        printHelpInfo();
        return EXIT_SUCCESS;
    }
    auto sampleTest=sample::gLogger.defineTest(gSampleName,argc,argv);
    sample::gLogger.reportTestStart(sampleTest);

    samplesCommon::CaffeSampleParams params = initializeSampleParams(args);
    SampleMNIST sample(params);
    sample::gLogInfo <<"Building and running a GPU inference engine for MNIST" << std::endl;
    if (!sample.build())
    {
        return sample::gLogger.reportFail(sampleTest);
    }

    if (!sample.infer())
    {
        return sample::gLogger.reportFail(sampleTest);
    }

    if (!sample.teardown())
    {
        return sample::gLogger.reportFail(sampleTest);
    }

    return sample::gLogger.reportPass(sampleTest);
}

cmake_minimum_required(VERSION 3.19 FATAL_ERROR)
project(perception)

#选择Debug模式
SET(CMAKE_BUILD_TYPE "Debug")

#选择Release模式
#SET(CMAKE_BUILD_TYPE "Release")

set(CMAKE_CXX_FLAGS   " -std=c++11 -fno-stack-protector" )

#Debug和Release模式优化级别
#set(CMAKE_CXX_FLAGS_DEBUG   "-O0 -g -fstack-protector -fstack-protector-all -lefence" )             # 调试包不优化
#set(CMAKE_CXX_FLAGS_DEBUG   "-O0 -g -fstack-protector-all -lefence" )             # 调试包不优化
#set(CMAKE_CXX_FLAGS_RELEASE "-O3 -DNDEBUG " )   # release包优化

set(CMAKE_CXX_FLAGS_DEBUG "$ENV{CXXFLAGS} -O0 -Wall -g2 -ggdb")
set(CMAKE_CXX_FLAGS_RELEASE "$ENV{CXXFLAGS} -O3 -Wall")

#设置可执行程序输出路径
set(EXECUTABLE_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/bin)
#设置库文件输出路径
set(LIBRARY_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/lib)

#使用相对路径,否则可执行程序和库文件更换路径后,程序无法运行
SET(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)
SET(CMAKE_INSTALL_RPATH "\${ORIGIN}/../lib") #指定运行时动态库的加载路径,ORIGIN指运行文件所在目录

find_package(CUDA REQUIRED)
# cuda
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# tensorrt
include_directories(/usr/include/x86_64-linux-gnu/)
link_directories(/usr/lib/x86_64-linux-gnu/)
find_package(OpenCV REQUIRED)

#find_package(Boost REQUIRED)
#find_package(Threads)
#include_directories(${Boost_INCLUDE_DIR})
#ADD_DEFINITIONS(-DBOOST_LOG_DYN_LINK)


#add_subdirectory(common)
#add_subdirectory(lidar)
set(DEOS_LIST "")
list(APPEND DEPS_LIST nvcaffeparser)
#target_include_directories(${TARGET_NAME}
#	PRIVATE common)

add_executable (${PROJECT_NAME}
                samplemnist.cpp
                common/logger.cpp)

target_link_libraries (${PROJECT_NAME}
	nvinfer
	nvcaffe_parser
	cudart
       	${OpenCV_LIBS}
	-lrt)

#redefine_file_macro(${PROJECT_NAME})