[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训练自己的数据集》

注意：

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

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

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

3. jit trace

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

cd yolov5

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

 import torchimport torchvisionimport modelsfrom models.experimental import attempt_loadtorch_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')

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

具体做法。修改model/yolo.py中Detect的forward函数，屏蔽掉条件判断处理部分。
执行trace。
```
 python3 jit_trace.py
```
trace成功后，会看到trace后的yolov5_jit.pt模型。

[Yolov5][Pytorch] 如何jit trace yolov5模型相关推荐

torch.jit.trace与torch.jit.script的区别
文章目录术语什么时候用torch.jit.trace(结论:首选) 优点什么时候用torch.jit.script(结论:必要时) 错误举例动态控制输入和输出有丰富类型的模型需要格外注意 Q ...
PyTorch导出JIT模型并用C++ API libtorch调用
PyTorch导出JIT模型并用C++ API libtorch调用本文将介绍如何将一个 PyTorch 模型导出为 JIT 模型并用 PyTorch 的 C++API libtorch运行这个模型 ...
[PyTorch] jit.script 与 jit.trace
pytorch 导出模型有两种方式,一种是 torch.jit.script,一种是 torch.jit.trace jit.script 如果网络中用到了索引,if 条件判断等,很有可能会失败. ...
海思 YOLOv5 pytorch 转 onnx 转 Caffe 再转 wk 的转化详解
目录: 前沿 YOLOv5模型的选取与修改 YOLOv5 pytorch 转 onnx 转 Caffe YOLOv5 Caffe转wk文件总结参考前沿作者在将YOLOv5 pytorch版本转 ...
Yolov5 系列1--- Yolo发展史以及Yolov5模型详解
最近在做检测相关的工作,几年前分析过faster-rcnn的代码和论文.现在,把yolov5这个最新且快的模型进行梳理.本系列会从Yolo的发展历程开始,到损失函数,mAP的概念,最后到如何在代码层面 ...
Yolov5自定义图片训练测试及模型调优（详细过程）
Yolov5自定义图片训练测试及模型调优(详细过程) 1. 图像数据获取.标注 2. 数据集划分及准备 3. 配置训练参数,准备训练: 3.1 修改data/训练配置xxx.yaml文件: 3.2 修 ...
基于改进 YOLOv5 的航空发动机表面缺陷检测模型如何制作？
建立基于改进 YOLOv5 的航空发动机表面缺陷检测模型主要需要以下步骤: 收集航空发动机表面缺陷的数据集.这些数据可以包括训练图像和标签数据,其中标签数据包含了航空发动机表面缺陷的位置信息. 利用 ...
C++调用Python文件，TensorFlow和PyTorch构建的深度学习模型，无法使用GPU的情况分析。
C++调用Python深度学习模型,包含TensorFlow和PyTorch等构造的模型,然后使用GPU出现问题.包含C++调用Python函数,C++加载模型到GPU,GPU内存占用过大,计算完毕内 ...
基于Pytorch和RDKit建立QSAR模型
尝试使用pytorch和RDKit构建QSAR模型环境依赖 pip install pprint pip install argparse #安装rdkit conda install -c rdk ...

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

1. 环境要求

2. 准备工作

3. jit trace

[Yolov5][Pytorch] 如何jit trace yolov5模型相关推荐

最新文章

热门文章