YOLO-cls训练及踩坑记录

提示：文章写完后，目录可以主动生成，如何生成可参考右边的帮助文档
  
   文章目录

  前言
  一、模子训练
  二、测试
  三、踩坑记录
  1、推理时设置的imgsz不生效
  方法一：
  方法二：
  2、Windows下torchvision版本问题导致报错
  总结
  

---

前言

        近期，想尝试下yolov11-cls的实际效果，以是在网上随便找了个花分类的数据集测试了下。数据准备阶段可以参考我的另一篇博客举行处理，链接如下：数据准备参考链接。接下来的正文部分就是我的训练和踩坑记录。留意:假造环境里不要安装ultralytics库，否则改动的代码不会生效！！！

---

一、模子训练

代码链接为：yolov11        
它的训练方式非常简单，查考yolov11官方文档即可，我的训练方法为：

1. from ultralytics import YOLO
2. model= YOLO("./yolo11n-cls.pt")
3. model.train(data="./datasets",epochs=100,imgsz=[256,256],batch=32)

复制代码

其中，模子路径和数据集路径需要改为自己的，epochs、imgsz、batch等根据自己的呆板环境设定。留意，数据集要根据我开头给的谁人链接里的格式摆放，即在datasets下包罗两级目录，如图：

然后，train和val下有每个类别的训练和测试数据。训练完成后，我们的模子默认存在./runs/classify下，如我的：

二、测试

我的测试代码如下：

1. from ultralytics import YOLO
3. model=YOLO('./runs/classify/train/weights/best.pt')
5. res=model("./test.jpg",imgsz=256,device="cuda")
6. print(res)

复制代码

Linux下根本可以正常执行，其他涉及到的问题继续往下看

三、踩坑记录

1、推理时设置的imgsz不生效

用如上代码推理时，发现日志输出包罗这个内容：

what？？？我设置的imgsz是256，为毛推理时还是224？估计是推理时尺寸强制用了训练时候的尺寸，假如你想在推理时更改尺寸，办理办法如下：
方法一：

根据这个issues的方法，在./ultralytics/engine/predictor.py里的setup_source函数(大概190行左右)举行更改，如下：
源代码：

更改后：

然后在执行推理即可。
方法二：

看到./ultralytics/models/yolo/classify/predict.py里的如下代码(留意红框里的代码)：

由红框里的代码可以知道，当你的img是非tensor时才会去举行预处理，以是你提前给它处理成tensor不举行if下的数据操纵即可，代码如下（预处理代码均是从官方预处理代码里拷贝的）：

完备代码可以我提的issues里获取。
2、Windows下torchvision版本问题导致报错

和服务器同样的torch和torchvision版本，但是在我自己的windows下推理时，抛出了如下异常：

这个问题就是torchvision版本的问题，经过多次尝试，发现如下版本可以正常执行。另外该报错也可通过问题1的办理方案办理。

1. pip install torch==1.8.1+cpu torchvision==0.9.1+cpu torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html

复制代码

留意，上边那条命令安装的是cpu的，假如你要装gpu的，去官方文档里找对对应的命令即可。

---

总结

以上就是本篇的全部内容，如有问题，欢迎评论区交流，或+企鹅群：995760755交流；如以为有用，欢迎三连

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。