在Xavier上运行yolov5，并解决pytorch版本兼容问题

本文主要描述我在Jetson AGX Xavier上运行u版yolov5:v4.0

我的Xavier配置：

• JetPack：4.3
• cuda:10.0
• python:3.6.9

YOLOv5

这里的下载训练推理均在服务器上运行。

下载yolov5：

git clone https://github.com/Oswells/yolov5.git

说明：我这里是fork了ultralytics/**yolov5**这个大神4.0版本的yolov5，目前已经推出5.0版本，下载最新的还是需要在这位大神的github上下载。

这里我主要说一下几个重要文件。

• models/yolo5s.yaml

  这里一共有四个类似文件，分别是s,m,l,x，模型依次增大。这四个文件定义了四个网络结构。如果对网络结构不关心，那么只需关注nc，你有几个类就改成几，我只有一个类，于是我设为1.

• data/hyp.scratch.yaml

  这个文件定义了一些超参数，训练前改动此文件以选用不同的超参数。从hsv_h开始都是数据增强相关的超参数，设为0时表示不启用。

训练yolov5

主要参考这篇博客，没错，作者提到的的热心网友就是我。

python选择3.7版本及以上，然后安装对应库，yolov5作者写了requirements.txt可以帮助我们很快部署环境。

pip3 install -r requirements.txt

作者也给了Dockerfile，如果有docker可以选择使用Dockerfile产生容器部署环境，我试过这个方法，可以用，但没必要hh。

yolov5使用了argparse库，这个库似乎已经成为神经网络程序的标配了，它可以灵活地从终端向程序传入参数。yolov5在train.py和detect.py中定义了非常多的可传入参数，但是有些参数一般我们都是固定某个值的，比如data（数据集所在位置），我们可以在程序中直接修改default值，这样就不用每次都在终端中敲同样的东西了。

yolov5推理

一般就运行以下代码就可以推理。

python detect.py --weights 'weights/best.pt' --source '1' --conf-thres 0.4 --device 'cpu'

这里weights是训练好的权重，source当它是数字的时候表示从摄像头读取图像，当它为文件夹路径时，会将整个文件夹里的图片/视频进行推理，当它为文件时，则推理该文件。conf-thres是置信度阈值，设置较低时容易将其他物体识别成目标。device选择程序在什么设备上运行，为cpu时运行较慢，亲测推理640x480图像只有2FPS，为数字时选择在显卡（gpu）运行，选择多张显卡时参考此格式’0,1’。

推理不需要那么多的文件，我们可以删减一些文件，这里我罗列推理需要用到的文件。

文件夹：models，utils，data

文件：权重文件，detect.py，

Xavier Pytorch 跑 yolov5

基本环境安装

建议安装miniforge，这个是anaconda的替代品，Xavier上无法安装anaconda。安装方式参考这篇博客.miniforge的使用方式与anaconda完全相同。

安装好后新建环境，python版本选择3.6，打开yolov5文件中的requirements.txt，除了opencv-python，torch和torchvision，一个个都用conda install安装，conda用不了再用pip3。opencv-python这个库xavier自带，不用安装。

torch和torchvision安装

由于Xavier是arm架构，所以直接pip或conda安装的torch和torchvision没办法使用，需要从这个网站下载对应版本的torch和torchvision。由于我的Xavier的Jetpack版本为4.3，所以没办法安装高版本的torch，这里我下载安装了1.4.0版本的torch和0.5.0版本的torchvision。

torch安装

pip3 install torch-1.4.0-cp36-cp36m-linux_aarch64.whl	#这里我默认whl文件在终端打开的目录中

torchvision安装

https://github.com/pytorch/vision/tags?after=v0.7.0-rc3 这个界面下载对应版本的torchvision，然后解压缩，cd进这个文件夹。

sudo python3 setup.py install 

解决兼容性问题

由于我们训练的时候用的是高版本的torch，实际使用的时候是1.4版本的torch，所以不可避免地会遇上兼容性问题。

• 更改torch模型保存格式

  首先保证自己电脑torch版本≥1.7.。

  在自己电脑上新建py文件，内容如下，自行修改load_file和load_path，load_path。

import torch
import os
load_file = "5s_multi_scale_gaussian"
load_path = os.path.join(load_file,"weights","best.pt")
load_path = os.path.join("weights",load_file+".pt")

ckpt = torch.load(load_path)
print(ckpt.get('model'))
ckpt = {'epoch': ckpt.get('epoch'),
        'best_fitness': ckpt.get('best_fitness'),
        'training_results': ckpt.get('training_results'),
        'model': ckpt.get('model'),
        'ema': ckpt.get('ema'),
        'updates': ckpt.get('updates'),
        'optimizer': ckpt.get('optimizer'),
        'wandb_id': ckpt.get('wandb_id')}
torch.save(ckpt,save_path,_use_new_zipfile_serialization=False)

运行程序，将新的权重文件copy到Xavier上。

注意，由于torch中使用Pickle保存加载模型，因此在必须要与原yolov5目录相似的目录结构才能成功load权重，这里必须要有models和utils文件夹与此py文件在同一目录下。

• 更改torch底层代码

  由于yolov5调用了新版torch才有的类，我们在使用老版torch时需要将这些函数和类补上。

  首先在Xavier中找到torch所在目录，在终端输入

  pip3 show torch
  

  在Location中可以找到它的存储位置，打开这个目录，再打开/nn/modules/activation.py，加入以下内容

  class SiLU(Module):
      def forward(self,x):
          return x * torch.sigmoid(x)
  
  
  class Hardswish(Module):
      def forward(self,x):
          return x * F.hardtanh(x + 3, 0., 6.) / 6.
  

• 修改yolov5/models/experimental.py

  开头加上from torch.nn.modules.activation import SiLU,Hardswish

  定位117行，如下修改

  #原先
  if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
  #修改为
  if type(m) in [Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, SiLU]:
  

  我觉得这里是我不懂得python底层的缘故，我不知道怎样让nn.SiLU和nn.Hardswish也能使用，我这种方法可能还是比较拙劣。

OK，现在就可以在终端中运行detect.py了。

tensorRT加速

我使用上面方法，在实际使用的时候发现yolo5s只有19FPS左右，达不到实时的效果，而听说tensorRT可以加速推理，所以我试着搞了一下，效果还不错，yolov5s最后达到了40FPS左右。

yolo转tensorRT

yolo的pt文件不能直接使用tensorRT加速，需要转换格式。github上有大神开发了一步到位的格式转换。这里我贴出转换yolov5::v4.0的链接：https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v4.0/yolov5

他的readme中Config第一点选择模型可以不用去管，其他内容自己按照他的说明设置。

设置完成后按照How to Run, yolov5s as example去做，一直做到第二点的sudo ./yolov5 -s yolov5s.wts yolov5s.engine s即可，这里由于我电脑没有gpu，所以我在xavier上执行上述步骤。到这里，我们就构造好了运行所需的权值文件.engine。

调用摄像头

我这里主要参考了这篇博客的第四点视频检测。由于他没有使用torch，所以我直接在base环境中使用engine。首先安装pycuda

pip3 install pycuda

然后新建py文件，将作者的代码copy上去，将关键参数修改即可，但是我在使用的过程中觉得有些可以改进的地方，我这里将修改后的函数写在下面。

def draw_boxes(image_raw, result_boxes,color, result_scores, result_classid):
    for i in range(len(result_boxes)):
        box = result_boxes[i]
        plot_one_box(
            box,
            image_raw,
            color,
            label="{}:{:.2f}".format(
                categories[int(result_classid[i])], result_scores[i]
            ),
        )
    return image_raw

def detect_camera(camera, yolov5_wrapper):
    ##开始循环检测，并将结果写到result.mp4中
    color = [random.randint(0, 255) for _ in range(3)]
    while True:
        #ret,img = camera.read()  # usb摄像头用这个
        t1 = time.time()
        img = camera.read()
        img, result_boxes, result_scores, result_classid = yolov5_wrapper.infer(img)
        img = draw_boxes(img, result_boxes,color, result_scores, result_classid)
        cv2.namedWindow("result", 0)
        cv2.resizeWindow("result", 640, 480)
        cv2.imshow("result", img)
        print("%d FPS"%(1/(time.time()-t1)))
        if cv2.waitKey(1) == ord('q'):
            break

注意修改main_camera()中的路径和usb摄像头。

# 没有调用CSICamera()方法则无需import
def main_camera():
    camera = cv2.VideoCapture(0)  # usb摄像头用这个，也可以用jetcam中的usb接口
    # camera = CSICamera(capture_device=0, width=224, height=224)
    # load custom plugins
    PLUGIN_LIBRARY = "build/libmyplugins.so"
    ctypes.CDLL(PLUGIN_LIBRARY)
    engine_file_path = "build/yolov5s.engine"

    # YoLov5TRT instance
    yolov5_wrapper = YoLov5TRT(engine_file_path)
    print("start detection!")
    detect_camera(camera, yolov5_wrapper)
    camera.release()  # 使用cv方法打开摄像头才需要这句
    cv2.destroyAllWindows()
    print("\nfinish!")
    
main_camera()

最后运行该py文件即可。