将 yolov5 内置模型导出为 pytorch mobile 格式

Google Colab 地址

使用 Colab 能利用 Google 免费提供的 GPU 资源，快速进行模型训练。比在我的弱鸡笔记本上训练，能节省大量时间。

在 Colab 上新建一个项目

https://colab.research.google.com

或者可以考虑继续沿用之前的项目。

安装 yolov5

新建一个代码块

#clone YOLOv5 and
!git clone https://github.com/ultralytics/yolov5  # clone repo
%cd yolov5
%pip install -qr requirements.txt # install dependencies
%pip install -q roboflow

import torch
import os
from IPython.display import Image, clear_output  # to display images

print(f"Setup complete. Using torch {torch.__version__} ({torch.cuda.get_device_properties(0).name if torch.cuda.is_available() else 'CPU'})")

点击运行，需要等待虚拟机的分配。

测试是否安装成功

用 yolo 内置的测试图片，判断是否能识别出人像。

!python detect.py --weights yolov5s.pt --source ./data/images/zidane.jpg
Image("runs/detect/exp/zidane.jpg", width=400)

如果能识别出齐祖和安猪，说明 yolov5 已经安装成功。

内置的模型

因为内置的模型能识别 80 种物体，有时候还是能直接使用的。

这个模型默认在 yolov5 的根目录下。

导出 pytorch mobile 模型

!python export.py --weights ./yolov5s.pt --include torchscript --optimize

导出的模型文件名为 yolov5s.torchscript 文件，在 .pt 文件的同级目录下。

tags: pytorch yolov5

关于作者 🌱

我是来自山东烟台的一名开发者，有敢兴趣的话题，或者软件开发需求，欢迎加微信 zhongwei 聊聊， 查看更多联系方式