在Ubuntu18上用LibTorch（非darknet）部署YOLOv4的填坑之路（附源码）

目录

坑一：从官网下载的LibTorch库是不带torchvision的

坑二：Python的PIL库与opencv库在图像处理上的差异值得注意

坑三：LibTorch对tensor的各种变换操作度相比Python令人窒息

坑四：LibTorch中的tensor转数组（向量）

坑五：YOLOv4模型的输出是个tuple，不能在forward后直接使用toTensor()

---

三周前，满怀懵逼的心情开始了艰难的YOLOv4的部署之路，有多艰难？一没基础，二没支援，刚开始用PyTorch，LibTorch从没接触过，甚至C++也要边学边做，唯一能问的就是小度和谷大哥。。。在某hub上找了好久有关YOLOv4的部署代码，竟然没发现用纯LibTorch来做的（也可能是我没找到吧），上面好多是用LibTorch+DarkNet写的，因为不想用DarkNet，就索性自己撸一个纯LibTorch的后处理代码吧。昨天测试之后终于初见成效，所以就把这部署过程中遇到的坑在此一一记录下来。

首先声明操作环境：

CPU版：Ubuntu18、PyTorch1.5.0CPU、LibTorch1.5.0CPU、CMake、VSCode、OpenCV 3.4.10、torchvision-0.6.0

GPU版：Ubuntu18、PyTorch1.7.1GPU、LibTorch1.7.1GPU、CMake、VSCode、CUDA10.2、torchvision0.8.2、OpenCV3.4.10

最终效果：LibTorch的目标输出数据与PyTorch的输出一致

代码我已经上传到GitHub，链接：

CPU版：https://github.com/wsx000/YOLOv4-LibTorch

GPU版：https://github.com/wsx000/YOLOv4-LibTorch-GPU

有关Ubuntu18下编译安装OpenCV的教程可以看这里。

坑一：从官网下载的LibTorch库是不带torchvision的

如果想用torchvision里面的函数，比如nms_cpu()，就需要自行下载并编译torchvision的。具体方法可以看这篇博客。

坑二：Python的PIL库与opencv库在图像处理上的差异值得注意

在用libtorch编完后处理代码后，除去代码本身因疏忽而导致的错误结果，发现跟PyTorch里跑出来的结果竟不一样！经过多番查证，发现就是PyTorch中使用PIL库处理图像的格式不同与opencv库，具体的不同如下：

1、图片对尺寸的存储格式上

使用PIL库直接读取图片的尺寸输出的是（w, h）。转换为numpy数组后就是（h, w, c）了，在用不同方式获取图像宽高时要特别注意。

opencv对图像的存储格式是（h, w, c）

可以看个例子，所用的图片宽500，高333：

from PIL import Image
import numpy as np
import cv2
pil_img = Image.open("/home/wsx/code/YOLOv4/y4-libtorch-V1-20201120/62.jpg")  # PIL库读取
cv_img = cv2.imread("/home/wsx/code/YOLOv4/y4-libtorch-V1-20201120/62.jpg")  # opencv读取
print("pil_img   :", pil_img.size)   # 查看PIL库读取的图像的size
print("np-pil_img:", np.shape(pil_img))  # 将PIL库读取的图像转换为numpy格式后打印shape
print("cv_img    :", cv_img.shape)   # opencv库读取图像的shape

#=================运行结果=================#
pil_img   : (500, 333)
np-pil_img: (333, 500, 3)
cv_img    : (333, 500, 3)

2、图片色彩通道的存储格式上

opencv对图片的存储格式是BGR，而PIL库则是RGB，所以如果PyTorch中使用PIL库处理图像，则LibTorch下用opencv处理的图像要转换为RGB格式才能是最终的输出结果一致。另外，不同库的同种差值方法计算的结果也不完全相同。

就以YOLO中无畸变调整图片的函数为例，说明两种方式的不同：

# PIL方法调整图片大小
def letterbox_image(image, size):
    iw, ih = image.size
    w, h = size
    scale = min(w/iw, h/ih)
    nw = int(iw*scale)
    nh = int(ih*scale)
    image = image.resize((nw,nh), Image.BICUBIC)
    new_image = Image.new('RGB', size, (128,128,128))
    new_image.paste(image, ((w-nw)//2, (h-nh)//2))
    return new_image

# opencv方法调整图片大小
def letterbox_image(image, size):
    new_image = np.ones((608,608,3),np.uint8)
    new_image[:] = 128
    h, w = image.shape[0], image.shape[1]
    scale = min(608./w, 608./h)
    nw, nh = int(scale*w), int(scale*h)
    image = cv2.resize(image,(nw,nh))
    offsetw, offseth = int((608-nw)/2), int((608-nh)/2)
    new_image[offseth:nh+offseth, offsetw:nw+offsetw] = image
    # cv2.imshow('hhh', new_image)
    # cv2.waitKey()
    new_image = cv2.cvtColor(new_image, cv2.COLOR_BGR2RGB)
    return new_image

坑三：LibTorch对tensor的各种变换操作度相比Python令人窒息

对YOLOv4模型的输出做后处理难免会遇到各种矩阵变换操作，PyTorch能做到的LibTorch都能做到，只是实现方式稍微复杂了一点点而已，在部署过程中所用到的所有LibTorch对矩阵的操作方法可以参考这篇博客和这篇。注意对矩阵的操作编程时一定仔细仔细再仔细！！！不然排查起来真令人头大！血泪教训啊啊啊(T...T)

坑四：LibTorch中的tensor转数组（向量）

最后总是要将tensor转换为数组来操作的，网上找了下，貌似很少有写如何转换的，下面放上我转换的方法，可能有些复杂，如果有更好的方法可以告诉我。

我是一个n*6的tensor转向量：

    // tensor转换为数组 (n, 6)  bboxes就是要转换的tensor，boxes是转换后的向量
    vector<vector<float>> boxes(bboxes.sizes()[0], vector<float>(6));

    for (int i = 0; i < bboxes.sizes()[0]; i++)
    {
        for (int j = 0; j < 6; j++)
        {
            boxes[i][j] = bboxes.index({at::tensor(i).toType(at::kLong),at::tensor(j).toType(at::kLong)}).item().toFloat();
            cout << boxes[i][j] << endl;
        }
        
    }

坑五：YOLOv4模型的输出是个tuple，不能在forward后直接使用toTensor()

这个坑其实是我遇到的第一个坑，当时直接在forward后面用了toTensor()方法就一直报错，还以为是模型转换的问题，后来一脸懵逼地排查了两天才发现问题(T...T)，正确的操作方法如下：

    // 前向传播
    auto outputs = module.forward(input).toTuple();

    // 提取三个特征层的输出
    vector<at::Tensor> out(3);
    out[0] = outputs->elements()[0].toTensor();
    out[1] = outputs->elements()[1].toTensor();
    out[2] = outputs->elements()[2].toTensor();

还有遇到其他很多小问题，解决方法就不在一一记录啦，作为新手难免有不足之处，有问题欢迎留言交流~