使用YOLOv8-gpu训练自己的数据集并预测
本篇将教学使用示例代码训练自己的数据集（train）以及预测（predict）。
目录
一、代码获取
二、训练
1、添加自己的训练集
2、创建训练集设置文件
3、 修改训练代码中数据集设置文件
4、开始训练
三、预测
1、 修改图片路径
2、 开始预测
总结
一、代码获取
在抠头助手安装YOLO系列库时都会自动下载示例代码，抠头助手安装YOLO系列的教程参考博客YOLOv5、v8、v11小白安装教程（GPU版和CPU版）_yolov11 cpu-CSDN博客，如果不需要装库，只需要下载代码可以直接前往训练代码进行下载（资源免费，点击下载即可）。
打开项目后文件结构如下：
- 橙色是预测示例图片
- 蓝色是训练集设置文件
- 红色是训练和预测代码
- 黄色是yolo模型文件
📦 项目根目录
├─ 📁 C包
├─ 📁 github
├─ 📁 .idea
├─ 📁 docker
├─ 📁 docs
├─ 📁 examples
├─ 📁 imgs
├─ 📁 my_training_results
├─ 📁 runs
├─ 📁 tests
├─ 📁 ultralytics
├─ .gitignore
├─ 说明书.md
├─ bus.jpg
├─ CITATION.cff
├─ coco128-seg.yaml
├─ coco128.yaml
├─ CONTRIBUTING.md
├─ cusPost.py
├─ cusPredict.py
├─ cusSeg.py
├─ cusTrain.py
├─ cusTrainSeg.py
├─ LICENSE
├─ mkdocs.yml
├─ pyproject.toml
├─ README.md
├─ README.zh-CN.md
├─ requirementsV8.txt
├─ yolov8n-cls.pt
├─ yolov8n-pose.pt
├─ yolov8n-seg-custom.pt
├─ yolov8n-seg.pt
├─ yolov8n.pt
└─ yolov8s.pt
如果需要其他模型可以前往GitHub - ultralytics/ultralytics at v8.2.103（需要魔法上网）下划找到如图所示位置下载模型，仍然按照上面的文件结构放置模型文件即可。
二、训练
1、添加自己的训练集
这里以“小目标检测数据集”为例。
- 解压数据集压缩包到代码目录下，位置如图所示：
2、创建训练集设置文件
- 创建该数据集的yaml设置文件，格式如下，根据自己数据集的类别和文件路径进行修改，比如：数据集的文件路径为D:\CodeTou-Download\YOLOV8_67\YOLOV8\小目标检测数据集，则path: D:\CodeTou-Download\YOLOV8_67\YOLOV8\小目标检测数据集；查看小目标数据集下的class.txt，有两个类别，则nc: 2；类别名称分别是earth_person，sea_person，则names: [ 'earth_person', 'sea_person' ] 。
# 小目标检测数据集配置文件
# 基于YOLO格式的小目标检测数据集# 数据集根目录（相对于yaml文件位置）
path: D:\CodeTou-Download\YOLOV8_67\YOLOV8\小目标检测数据集 # dataset root dir# 训练集和验证集路径
train: images/train # train images (relative to 'path')
val: images/val # val images (relative to 'path')
test: images/test # test images (optional)# 类别数量
nc: 2# 类别名称列表
names: ['earth_person', 'sea_person']
3、 修改训练代码中数据集设置文件
这里以检测模型为例，使用cusTrain.py来训练。
- 如果训练终止出现Memory Error，一般是内存占用满导致的，可以修改batch大小。
- 如果想要训练效果更好可以适当增加epochs
4、开始训练
- IDE鼠标点击运行
- IDE终端代码运行
python cusTrain.py
5、 训练结果
- 训练结果在runs/目录下查看。
三、预测
这里使用cusPredict.py
1、 修改图片路径
在第11行处修改要预测的图片路径。
2、 开始预测
- IDE鼠标点击运行
- IDE终端代码运行
python cusPredict.py总结
以上就是使用示例代码训练和预测的流程。