1. 下载YOLOv5官方代码
ONNX > CoreML > TFLite (github.com)">ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite (github.com)
2. 配置环境
打开终端,先建立名为YOLO5的环境,再将路径切换为requirements.txt文件夹所在的路径,再输入以下命令
pip install -r requirements.txt
操作没有问题的情况下,等待片刻环境就配置成功了。
3. 准备数据集
由于这是分类任务,数据集是不用标注的。只需先划分训练集、测试集、验证集,再将对应类别的图片放到对应类别的文件夹下即可。
这里的话每种类别400张图片,在训练集中有320张,测试集中有40张,验证集中有40张。
4. 准备预训练模型
如果让代码自动下载,会下载很久,推荐预先下载好了之后放在主文件夹下。
你可以访问以下链接下载预训练模型yolov5s-cls.pt:
https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s-cls.pt
如图所示
5. 训练模型
由于是分类任务,训练时用到的是classify文件夹下的train.py。数据会用到训练集和测试集。
找到parse_opt函数
,设置命令行参数
parser.add_argument("--model", type=str, default="yolov5s-cls.pt", help="initial weights path")
设置预训练模型为yolov5s-cls.pt
parser.add_argument("--data",