text: #include "NvInfer.h"
  #include "cuda_runtime_api.h"
  #include "logging.h"
  #include <fstream>
  #include <map>
  #include <chrono>
  
  #define CHECK(status) \
      do\
      {\
          auto ret = (status);\
          if (ret != 0)\
          {\
              std::cerr << "Cuda failure: " << ret << std::endl;\
              abort();\
          }\
      } while (0)
  
  // stuff we know about the network and the input/output blobs
  static const int INPUT_H = 32;
  static const int INPUT_W = 32;
  static const int OUTPUT_SIZE = 10;
  
  const char* INPUT_BLOB_NAME = "data";
  const char* OUTPUT_BOLB_NAME = "prob";
  
  using namespace nvinfer1;
  
  static Logger gLogger;
  
  // Load weights from files shared with TensorRT samples.
  // TensorRT weight files have a simple space delimited format:
  // [type] [size] <data x size in hex>
  std::map<std::string, Weights> loadWeights(const std::string file)
  {
      std::cout << "Loading weights: " << file << std::endl;
      std::map<std::string, Weights> weightMap;
  
      // Open weights file
      std::ifstream input(file);
      assert(input.is_open() && "Unable to load weight file.");
  
      // Read number of weight blobs
      int32_t count;
      input >> count;
      assert(count > 0 && "Invalid weight map file.");
  
      while (count--)
      {
          Weights wt{DataType::kFLOAT, nullptr, 0};
          uint32_t size;
  
          // Read name and type of blob
          std::string name;
          input >> name >> std::dec >> size;
          wt.type = DataType::kFLOAT;
  
          // Load blob
          uint32_t* val = reinterpret_cast<uint32_t*>(malloc(size(val)*size));
          for (uint32_t x = 0, y = size; x < y; ++x)
          {
              input >> std::hex >> val[x];
          }
          wt.values = val;
  
          wt.count = size;
          weightMap[name] = wt;
      }
  
      return weightMap;
  }
  
  // Create the engine using only the API and not any parser.
  ICudaEngine* createLenetEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt)
  {
      INetworkDefinition* network = builder->createNetworkV2(0U);
  
      // Create input tensor of shape { 1, 32, 32 } with name INPUT_BLOB_NAME
      ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{1, INPUT_H, INPUT_W});
      assert(data);
  
      // Add convolution layer with 6 outputs and a 5x5 fiter.
      std::map<std::string, Weights> weightMap = loadWeights("../lenet5.wts");
      IConvolutionLayer* conv1 = network->addConvolutionNd(*data, 6, DimsHW{5, 5}, weightMap["conv1.weight"], weightMap["conv1.bias"]);
      assert(conv1);
      conv1->setStrideNd(DimsHW{1, 1});
  
      // Add activation layer using the ReLU algorithm
      IActivationLayer* relu1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU);
      assert(relu1);
  
      // Add max pooling layer with stride of 2x2 and kernel size of 2x2.
      IPoolingLayer* pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kAVERAGE, DimsHW{2, 2});
  }
  
  void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream)
  {
      // Create builder
      IBuilder* builder = createInferBuilder(gLogger);
      IBuilderConfig* config = build->createBuilderConfig();
  
      // Create model to populate the network, then set the outputs and create an engine
      ICudaEngine* engine = createLenetEngine(maxBatchSize, builder, config, DataType::kFLOAT);
      assert(engine != nullptr);
  
      // Serialize the engine
      (*modelStream) = engine->serialize();
  
      // Close everything down
      engine->destory();
      builder->destory();
  }
  
  void doInference(IExecutionContext& context, float* input, float* output, int batchSize)
  {}
  
  int main(int argc, char** argv)
  {
      // argc相关
      if (argc != 2) {
          std::cerr << "arguments not right!" << std::endl;
          std::cerr << "./lenet -s   // serialize model to plan file" << std::endl;
          std::cerr << "./lenet -d   // deserialize plan file and run inference" << std::endl;
          return -1;
      }
  
      // create a model using the API directly and serialize it to a stream
      char *trtModelStream{nullptr};  // 创建一个空指针trtModelStream
      size_t size{0};  //size_t size = 0;
      
      // 生成engine文件
      if (std::string(argv[1]) == "-s")
      {
          IHostMemory* modelStream{nullptr};  // 创建一个IHostMemory类型的指针
          APIToModel(1, &modelStream);
          assert(modelStream != nullptr);
  
          std::ofstream p("lenet5.engine", std::ios::binary);
          if (!p)
          {
              std::cerr << "could not open plan output file" << std::endl;
              return -1;
          }
          p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
          modelStream->destory();
          return 1;
      } else if (std::string(argv[1]) == "-d")
      {
          std::ifstream file("lenet5.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();
          }
      } else {
          return -1;
      }
  }