YOLOv11训练完毕后追加训练的方法(不改变学习率)
一位没玩过yolo的小白最近第一次使用yolo11训练模型,但是发现最初设定的150轮训练并没有达到一个好的训练效果,于是希望在原来基础上继续训练个150轮,也就是300轮。
如果你有同样的需求,可以往下看(我的yolo11版本8.3.86):
以下是我原来训练150轮的代码:
from ultralytics import YOLO
# 加载模型
model = YOLO('yolo11s.yaml') # 从 YAML 文件创建新模型
# 训练模型
results = model.train(
data='TrashCan.yaml', # 数据集配置文件
epochs=150 , # 训练周期数
batch=32, # 每批次图像数量
imgsz=256, # 输入图像尺寸
device=0, # 使用 GPU 0
workers=0,
project='runs/train', # 保存训练结果的目录
name='exp1', # 实验名称
)
询问ai后,说只要修改为以下代码就可以恢复训练:
from ultralytics import YOLO
# 加载上次保存的检查点文件
model = YOLO('runs/train/exp1/weights/last.pt')
# 训练模型
results = model.train(
data='TrashCan.yaml', # 数据集配置文件
epochs=300 , # 训练周期数
batch=32, # 每批次图像数量
imgsz=256, # 输入图像尺寸
device=0, # 使用 GPU 0
workers=0,
project='runs/train', # 保存训练结果的目录
name='exp1', # 实验名称
resume=True
)
但是一运行,就被告知:真不好意思,你已经训练150轮了,还想在训练,不可能!
也就是说,以上方法仅仅适用于意外中断或者主动中断还未训练完的模型才可以恢复训练。
于是我继续询问ai怎么解决,又被告知:既然不行,你别使用resume=True进行恢复训练了,直接在这个基础上再训练150轮吧!
但是,这样存在一个问题,就是学习率会被重置,导致训练一开始,因学习率过大而发散。
鉴于yolo11版本比较新,其他yolo版本的解决方案不适用,于是花了几个小时看了下源代码,找到了解决方案:
第一步:打开\ultralytics-main\runs\train\exp1文件夹下的args.yaml文件,将其中的epochs修改为你第二遍想要达到的轮数(包括第一遍已经训练完毕的轮数),我这里第一遍跑完了150个epoch,所以改成了epochs: 300,意思就是再跑150轮。
第二步:在\ultralytics-main目录下新建一个python代码,也可以随便找个地方建,路径自己改一下,代码内容如下:
import torch
import os
# 检查点文件路径
checkpoint_path = './runs/train/exp1/weights/last.pt'
# 加载检查点文件
ckpt = torch.load(checkpoint_path, weights_only=False)
# 修改epoch参数为149
ckpt['epoch'] = 149
# 修改epochs参数为300
ckpt['train_args']['epochs'] = 300
# 保存修改后的检查点文件
torch.save(ckpt, checkpoint_path)
print(f"成功将 {checkpoint_path} 中的epoch参数修改为149,epochs参数修改为300")
解释一下,该代码用来修改last.pt中的两个参数,第一个epoch,代表已经训练了的轮数减一,再强调一遍,是已经训练了的轮数减一(源代码里就是这么整的,我看过,你们别管为什么了),我这里因为第一遍跑了150轮,所以改成149。第二个参数epochs,代表训练总轮数,我是300。
第三步,把自己的训练代码里的epochs改成300,就大功告成了!!!
from ultralytics import YOLO
# 加载上次保存的检查点文件
model = YOLO('runs/train/exp1/weights/last.pt')
# 训练模型
results = model.train(
data='TrashCan.yaml', # 数据集配置文件
epochs=300 , # 训练周期数
batch=32, # 每批次图像数量
imgsz=256, # 输入图像尺寸
device=0, # 使用 GPU 0
workers=0,
project='runs/train', # 保存训练结果的目录
name='exp1', # 实验名称
resume=True
)
可以看到,代码愉快的继续跑起来了,而且学习率也保持原来的,各个指标都在收敛。