阅读量:2
1. 下载YOLOv5官方代码
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", 