快速上手YOLOv9目标检测算法：从配置到推理再到训练自定义模型

目录

1. 前言

2.下载安装NVIDIA驱动&CUDA

3.下载YOLOV9项目到本地

4.模型推理

5.训练自己的模型

6. 问题解决

1.GTX16系列的显卡训练出现loss值为nan

2.Could not run ' torchvision: :nms' with arguments from the ' CUDA' backend

3.其他问题等待总结 

---

1. 前言

YOLOv8 正式发布一年多以后，YOLOv9 终于问世，YOLO在目标检测领域中以其高速和高精度而闻名。本文将介绍如何在使用NVIDIA显卡的环境下进行YOLOv9模型的下载、安装、模型推理和自定义训练。并侧着介绍怎么使用roboflow对简单目标进行打自动标签。通过本文，您将学会如何配置环境、下载项目、进行模型推理和自定义训练，并解决可能出现的一些常见问题。以下是具体步骤：

2.下载安装NVIDIA驱动&CUDA

1. 下载安装NVIDIA驱动

   先检查是否安装了驱动

   nvidia-smi
   

2. 如果没有这个命令则代表没有安装NVIDIA驱动，或者你觉得NVIDIA驱动版本太落后，都可以前往Official Drivers | NVIDIA更新。选择默认的游戏版本(Game Ready Driver)即可，不需要选择工作室版本，那个是给图形设计用的。

3. 检查CUDA

   nvcc --version
   

4.  如果没有这个命令则没有安装cuda，请前往CUDA Toolkit 12.4 Downloads | NVIDIA Developer下载当前NVIDIA驱动支持的cuda(上面nvidia-smi命令中有显示)

3.下载YOLOV9项目到本地

• GitHub链接：GitHub - WongKinYiu/yolov9: Implementation of paper - YOLOv9: Learning What You Want to Learn Using Programmable Gradient Information

• 推荐自己先手动下载合适的pytorch依赖库，下载地址Start Locally | PyTorch，自行选择环境(pip或conda)进行下载。

• 安装依赖

  pip install -r requirements.txt
  

4.模型推理

1. 下载官方提供的模型gelan-c.pt，下载地址：https://github.com/WongKinYiu/yolov9/releases/download/v0.1/gelan-c.pt

2. 复制模型到项目根目录下

3. 更改detect.py，找到位置def parse_opt():

   parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'gelan-c.pt', help='model path or triton URL')
   parser.add_argument('--source', type=str, default=ROOT / '填上你的资源路径', help='file/dir/URL/glob/screen/0(webcam)')
   ​
   # store_false改为store_false可以预览推理后的图片
   parser.add_argument('--view-img', action='store_false', help='show results')
   

4. 运行detect.py查看效果

5. 如果这里有出现问题，请翻到下面问题解决看看能否解决你的问题

5.训练自己的模型

1. 准备数据集

   说明：这里笔者使用对某一游戏的屏幕截图作为数据集，首先编写截图脚本。

   # -*- coding: utf-8 -*-
   """
   @author 星辰
   @date 2024年03月20日 18:07:01
   @packageName 
   @className screenshoter
   @version 1.0.0
   @describe TODO
   """
   import os
   import pyautogui
   import keyboard
    
   # 创建一个新文件夹来存放截图
   folder_name = "screenshots"
   if not os.path.exists(folder_name):
       os.makedirs(folder_name)
    
   count = 0  # 初始化计数器
    
    
   def take_screenshot():
       global count
       count += 1
       screenshot_path = os.path.join(folder_name, f"{count}.png")
       pyautogui.screenshot(screenshot_path)
    
       print(f"截图已保存: {screenshot_path}")
    
    
   # 监听键盘事件，按下数字键9时触发截图操作
   keyboard.on_press_key('9', lambda _: take_screenshot())
    
   print('截图已准备就绪\n请按下数字键9进行截图\n截图后会保存到当前项目下的 screenshots 文件夹\n按下ESC键则退出截图')
   keyboard.wait('esc')
   

2. 使用roboflow自动标注

   1. 注册登录操作省略

   2. 创建一个项目

 

 

  

 

注意：免费用户上传的图片会发到公共空间，请不要上传隐私图片

3.下载导出的数据集到本地项目根目录

4.打开train.py ，跳到函数def parse_opt(known=False):

# 添加权重
parser.add_argument('--weights', type=str, default='gelan-c.pt', help='initial weights path')
# 添加模型配置
parser.add_argument('--cfg', type=str, default='models/detect/gelan-c.yaml', help='model.yaml path')
# 添加自定义的模型配置
parser.add_argument('--data', type=str, default=ROOT / 'gamedata/data.yaml', help='dataset.yaml path')
# 更改超参数配置文件
parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-high.yaml', help='hyperparameters path')
# 更改训练轮数
parser.add_argument('--epochs', type=int, default=100, help='total training epochs')
# 更改GPU显存大小，是多少就填多少，可以少填，不要多填
parser.add_argument('--batch-size', type=int, default=2, help='total batch size for all GPUs, -1 for autobatch')
 
# ...
 
# 更改线程数，多少核就填多少
parser.add_argument('--workers', type=int, default=4, help='max dataloader workers (per RANK in DDP mode)')

 5.点击运行开始训练

6. 问题解决

1.GTX16系列的显卡训练出现loss值为nan

在train.py中找到amp = check_amp(model)，直接让他amp = False

2.Could not run ' torchvision: :nms' with arguments from the ' CUDA' backend

出现错误：NotImplementedError: Could not run ' torchvision: :nms' with arguments from the ' CUDA' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective

解决：重新下载torch即可

1.卸载torch和torchvision

pip uninstall torch torchvision

2.下载torch，下载地址Start Locally | PyTorch

3.其他问题等待总结