[Yolov5][Pytorch] 如何jit trace yolov5模型

1. 环境要求

为了确保不会出现因环境问题导致的错误，强烈建议使用下面的环境：

torch：1.7.0
torchvision：0.7.0
python： 3.7
yolov5: release <= 3.1

2. 准备工作

Clone 代码和下载预训练权重，详细请参考这篇文章《YOLOv5的详细使用教程，以及使用yolov5训练自己的数据集》

注意：

1. yolov5 的代码需要checkout到对应的tag。

   如：git checkout v3.1

2. 预训练的模型也要下载对应版本。如v3.1 release。

   https://github.com/ultralytics/yolov5/releases/tag/v3.1

3. jit trace

下面以Yolov5s为例，描述trace的过程。

1. cd yolov5

2. 新建一个jit_trace.py的文件，代码如下：

    import torch
    import torchvision
   
    import models
    from models.experimental import attempt_load
   
    torch_model = attempt_load('weights/yolov5s.pt', map_location=torch.device('cpu'))
    torch.jit.trace(torch_model, torch.rand([1, 3, 640, 640])).save('yolov5_jit.pt')
   

3. 因为jit_trace不支持条件和循环，所以模型最后的后处理部分不能trace，需要屏蔽掉。这部分放在模型外通过代码来实现。

   具体做法。 修改model/yolo.py中Detect的forward函数，屏蔽掉条件判断处理部分。
   

4. 执行trace。

    python3 jit_trace.py
   

   trace成功后，会看到trace后的yolov5_jit.pt模型。