libtorch部署yolov5调用cuda

要在 libtorch 中部署 YOLOv5 并调用 CUDA，您需要执行以下步骤：

1. 安装 CUDA 和 CuDNN

在安装 libtorch 之前，您需要先安装 CUDA 和 CuDNN。确保您的 CUDA 版本与您的 libtorch 版本兼容。

2. 下载 YOLOv5 模型

在部署之前，您需要下载 YOLOv5 模型。您可以从 GitHub 上的 YOLOv5 仓库中下载预训练模型。

3. 加载模型

在 libtorch 中加载模型非常简单。您只需要使用 `torch::jit::load()` 函数将模型加载到内存中。例如：

torch::jit::script::Module module = torch::jit::load("path/to/model.pt");

4. 设置输入数据

在使用模型之前，您需要准备输入数据。对于 YOLOv5，输入数据应为一张图片。您可以使用 OpenCV 或其他库来加载图片，并将其转换为张量。例如：

cv::Mat image = cv::imread("path/to/image.jpg");
cv::Mat image_float;
image.convertTo(image_float, CV_32F, 1.0 / 255.0);
torch::Tensor input_tensor = torch::from_blob(image_float.data, {1, image_float.rows, image_float.cols, 3}).permute({0, 3, 1, 2});

5. 将模型和输入数据移动到 GPU 上

在调用模型之前，您需要将模型和输入数据移动到 GPU 上。例如：

module.to(torch::kCUDA);
input_tensor = input_tensor.to(torch::kCUDA);

6. 调用模型

调用模型非常简单。您只需要将输入数据传递给模型并获取输出。例如：

torch::Tensor output_tensor = module.forward({input_tensor}).toTensor();

7. 后处理

对于 YOLOv5，输出是一个包含检测框和类别的张量。您需要对输出进行后处理，以便将张量转换为检测框和类别。例如：

std::vector<std::vector<float>> detections;
for (int i = 0; i < output_tensor.size(0); i++) {
    std::vector<float> detection;
    detection.push_back(output_tensor[i][0].item<float>());
    detection.push_back(output_tensor[i][1].item<float>());
    detection.push_back(output_tensor[i][2].item<float>());
    detection.push_back(output_tensor[i][3].item<float>());
    detection.push_back(output_tensor[i][4].item<float>());
    detection.push_back(output_tensor[i][5].item<float>());
    detections.push_back(detection);
}

这是一个基本的 YOLOv5 模型部署流程，您可以根据您的需求进行修改和优化。