目标
- pip安装依赖
- (可选)安装GPU版pytorch
- 自定义训练配置,模型训练
步骤
文档
pip安装依赖
# 安装依赖
pip install ultralytics
# 耗时有点长,可以使用国内镜像安装
pip install -i https://mirrors.aliyun.com/pypi/simple/ ultralytics
- 测试效果
- 随便搜索下载一张猫或狗的图片,把路径作为source参数
- "yolov8n.pt"是官方预训练好的模型,第一次使用会自动从Github上下载
- 国内访问Github网络不稳定,如果失败多重试几次,或自己下载放到根目录中
yolo predict model=yolov8n.pt source='dog.jpg'
![](https://img2023.cnblogs.com/blog/1100429/202309/1100429-20230916194531028-2123997599.jpg)
(可选)重装GPU版pytorch
- ultralytics默认安装的是cpu版pytorch
- 导致使用英伟达显卡并安装了CUDA和cuDNN,但训练时依然是CPU
- 运行下面的python文件,查看打印内容就可以看出来
import torch
# 判断cuda是否生效
print(torch.cuda.is_available())
# 查看torch版本
print(torch.__version__)
- pip卸载CPU版本pytorch
- 打开pytorch官网,重新安装GPU版pytorch
# 卸载cpu版本的pytorch
pip uninstall torch torchvision torchaudio
# 重新安装gpu版本(cuda11.7)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117
- 重新运行上面的python文件,看到版本从"2.0.1+cpu"变为"2.0.1+cu117"
自定义训练配置
- 参考ultralytics项目的coco128.yaml训练配置,编写自定义训练配置
- 将训练集根目录设置为上一期视频,从标注平台导出的目录
- 并将train、val路径设置为images(标注档会自动识别为"labels")
- 示例如下:game.yml
path: E:\ai-play-game\project-1-at-2023-09-13-17-05-6275bec0 # dataset root dir
train: images # train images (relative to 'path') 128 images
val: images # val images (relative to 'path') 128 images
# Classes,最好使用英文,opencv支持中文比较麻烦
names:
0: hellhound # 地狱犬
1: samurais # 武士
2: player # 玩家
3: fireflies # 萤火虫
训练
- 详细训练参数看官网ultralytics - train
- data: 指定训练配置路径
- model: 使用"yolo8n.yaml"模型(定义模型网络,"yolov8n.yaml"是内置yolov8n网络结构)
- batch: 批大小,数值太大电脑扛不住
- epoches: 训练轮次,游戏元素比较简单,训练一两百次就有不错的效果了
yolo detect train data=game.yml model=yolov8n.yaml batch=4 epochs=200
训练结果验证
- 训练时,自动在根目录创建runs/detect/train目录,可以看到各种训练数据
- 训练结束后,"runs/detect/train/weights/"目录里是模型结果
- best.pt,是本轮训练中效果最好的模型
- last.pt,是最后一次训练的模型
- 使用任意游戏截图测试效果
yolo predict model=runs/detect/train/weights/best.pt source='E:/project-e8d01b35/images/xxx.jpg'
![](https://img2023.cnblogs.com/blog/1100429/202309/1100429-20230916194544409-1175937618.jpg)