提供一个c++端用libtorch和Torch TensorRT加速推理torchscript模型的案例并注释

以下是一个使用libtorch和Torch TensorRT加速推理torchscript模型的C++案例：

#include <torch/script.h>
#include <iostream>
#include <memory>
#include <chrono>

#include <NvInfer.h>
#include <NvInferRuntimeCommon.h>
#include <NvInferPlugin.h>
#include <NvInferPluginUtils.h>

using namespace std::chrono;

int main(int argc, const char* argv[]) {
  if (argc != 3) {
    std::cerr << "Usage: " << argv[0] << " <model_path> <input_shape>" << std::endl;
    return 1;
  }

  const std::string model_path = argv[1];
  const std::string input_shape_str = argv[2];
  const std::vector<int64_t> input_shape = parse_input_shape(input_shape_str);

  // Load the model
  std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(model_path);

  // Create the input tensor
  torch::Tensor input = torch::randn(input_shape);

  // Warm up the GPU
  module->forward({input}).toTensor();

  // Convert the model to TensorRT
  const int max_batch_size = 1;
  const int max_workspace_size = 1 << 30; // 1 GB
  const int max_dla_batch_size = 0;
  const int max_dla_core = -1;
  const bool fp16_mode = true;
  const bool int8_mode = false;
  const int int8_calibration_batch_size = 0;
  const std::string engine_path = "engine.trt";
  nvinfer1::ICudaEngine* engine = convert_to_tensorrt(module, input_shape, max_batch_size, max_workspace_size,
                                                      max_dla_batch_size, max_dla_core, fp16_mode, int8_mode,
                                                      int8_calibration_batch_size, engine_path);

  // Create execution context
  nvinfer1::IExecutionContext* context = engine->createExecutionContext();

  // Allocate the output tensor
  std::vector<int64_t> output_shape = get_output_shape(engine, input_shape);
  torch::Tensor output = torch::empty(output_shape);

  // Run inference using TensorRT
  auto start = high_resolution_clock::now();
  std::vector<void*> buffers(2);
  const int input_index = engine->getBindingIndex("input");
  const int output_index = engine->getBindingIndex("output");
  buffers[input_index] = input.data_ptr();
  buffers[output_index] = output.data_ptr();
  context->executeV2(buffers.data());
  auto stop = high_resolution_clock::now();
  auto duration = duration_cast<microseconds>(stop - start);
  std::cout << "Inference time: " << duration.count() << " microseconds" << std::endl;

  // Verify the output
  torch::Tensor expected_output = module->forward({input}).toTensor();
  assert(torch::allclose(output, expected_output, 1e-3, 1e-3));

  // Clean up
  context->destroy();
  engine->destroy();
  return 0;
}

代码中的 `parse_input_shape`、`convert_to_tensorrt`、`get_output_shape` 是一些辅助函数，用于解析输入形状，将模型转换为TensorRT引擎和获取输出形状。这些函数的实现可以参考官方文档。

使用时，需要在命令行中指定模型路径和输入形状，例如：`./my_app my_model.pt 1,3,224,224`。程序会自动将模型转换为TensorRT引擎，并运行推理。最后，程序会验证TensorRT输出是否与PyTorch输出一致，以确保结果正确。