1.源代码
原本代码的内容如下:
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO
model = YOLO("yolov8m.pt")
results = model.train(data="yolo-dataset/yolo.yaml", epochs=32, imgsz=1080, batch=16)
2.上述代码分析
上述使用了YOLOv8模型进行训练,代码中包含以下关键步骤:
1、设置CUDA设备:
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
这行代码指定了使用哪块GPU进行训练。"0"
表示使用第一块GPU。如果有多块GPU,你可以指定其他的编号。
要查看当前系统中有多少个可用的GPU,你可以使用Python中的torch
库。以下是使用torch
检查GPU数量的方法:
import torch
# 检查是否有可用的GPU
if torch.cuda.is_available():
print("CUDA is available!")
# 获取GPU的数量
num_gpus = torch.cuda.device_count()
print(f"Number of GPUs available: {num_gpus}")
else:
print("CUDA is not available.")
这段代码首先检查CUDA是否可用,如果可用,则打印出可用的GPU数量。
2、忽略警告:
warnings.filterwarnings('ignore')
这行代码会忽略所有的警告信息,确保输出的日志更简洁。
3、加载YOLOv8模型:
model = YOLO("yolov8m.pt")
使用了yolov8m.pt
模型,这是YOLOv8的中等版本(medium)。YOLOv8有多个版本,yolov8n.pt
(nano),yolov8s.pt
(small),yolov8m.pt
(medium),yolov8l.pt
(large),yolov8x.pt
(extra-large)。每个版本的参数量和计算需求不同。
4、训练模型:
results = model.train(data="yolo-dataset/yolo.yaml", epochs=64, imgsz=1088, batch=16)
在这个代码中设置了训练的参数:
data="yolo-dataset/yolo.yaml"
:定义了数据集的配置文件。epochs=64
:训练64个周期。imgsz=1088
:训练时使用的图像尺寸为1088x1088像素。batch=16
:每个训练批次的大小为16张图片。
3.提高精度的建议及对应代码
1、增加数据增强
数据增强是提高模型泛化能力的有效方法。你可以在yolo.yaml
配置文件中添加或增强数据增强技术,如随机旋转、裁剪、缩放等。
# yolo.yaml 文件中的部分内容
augment:
flipud: 0.5 # 随机上下翻转
fliplr: 0.5 # 随机左右翻转
hsv_h: 0.015 # HSV色调调整
hsv_s: 0.7 # HSV饱和度调整
hsv_v: 0.4 # HSV亮度调整
translate: 0.1 # 平移
scale: 0.5 # 缩放
shear: 0.01 # 剪切
perspective: 0.001 # 透视变换
2、调用多个GPU
如果你想使用所有可用的GPU,可以设置CUDA_VISIBLE_DEVICES
环境变量为空或者包含所有GPU的编号。例如,如果你有4个GPU,可以设置为0,1,2,3
。
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3" # 假设有4个GPU
3、使用更大的模型
如果你的GPU允许,可以考虑使用更大的模型(例如yolov8x.pt
)来提高精度。这会增加模型的容量,但也需要更多的计算资源。(注:我在使用4090GPU24G显存的情况下,跑yolov8m模型进行微调大约用了19.8G的显存,而yolov8l与yolov8x则需要远远高于m的显存,小心爆显存!)
model = YOLO("yolov8x.pt")
4、调整学习率调度器
学习率调度器可以帮助模型在训练过程中更好地收敛。YOLOv8内部已经包含了自动调整学习率的机制,但你可以手动设置lr0
(初始学习率)和lrf
(最终学习率)来微调。
results = model.train(data="yolo-dataset/yolo.yaml", epochs=64, imgsz=1088, batch=16, lr0=0.001, lrf=0.01)
5、增加正则化
可以通过设置weight_decay
来增加正则化,防止模型过拟合。
results = model.train(data="yolo-dataset/yolo.yaml", epochs=64, imgsz=1088, batch=16, weight_decay=0.0005)
6、使用提前停止
提前停止是一种防止过拟合的方法,当模型在验证集上不再改善时提前停止训练。
results = model.train(data="yolo-dataset/yolo.yaml", epochs=64, imgsz=1088, batch=16, patience=10)
4.综合改进后的代码
综合以上建议,改进后的代码可能如下:
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO
# 使用更大的模型
model = YOLO("yolov8x.pt")
# 训练模型,包含数据增强、学习率调整、正则化和提前停止
results = model.train(
data="yolo-dataset/yolo.yaml",
epochs=32,
imgsz=1088,
batch=16,
lr0=0.001,
lrf=0.01,
weight_decay=0.0005,
patience=10
)