Yolov8改进模型后使用预训练权重迁移学习训练自己的数据集_yolo迁移训练

windows系统下YoLov8改进模型训练自己的数据集

yolov8 github下载

1、此时确保自己的数据集格式是yolo 格式的（不会的去搜教程转下格式）。

你的自制数据集文件夹摆放

主目录文件夹摆放

自制数据集data.yaml文件路径模板

path: ../datasets/VOCdevkit 
train: images/train  # train images (relative to 'path') 128 images
val: images/val  # val images (relative to 'path') 128 images
test:  # test images (optional)

names: 
  0: name1
  1: name2
  .....
1
2
3
4
5
6
7
8
9

2、把data.yaml放在yolov8–>ultralytics->datasets文件夹下面

3、然后模型配置改进yaml文件在主目录新建文件夹v8_cfg存放

模型改进修改与yolov5 类似，只是修改的文件名称与文件地址变化了

1.首先在yolov8–>ultralytics–>nn–>model.py添加你自己的新模型

2.然后在yolov8–>ultralytics–>nn–>task.py–>parse_model解析文件里面导入并添加你的模块

3.修改yolov8n.yaml文件

4、权重.pt文件直接下载到主目录下面

5、然后在主目录下新建推理v8_train.py

from ultralytics import YOLO

if __name__ == '__main__':
    # 加载模型
    model = YOLO("./v8_cfg/yolov8n-test.yaml")  # build a new model from scratch
    results =model.train(data="data.yaml" ,epochs=500, model="yolov8n-test.yaml",imgsz=640,batch=30,workers = 2)
1
2
3
4
5
6

此时运行改进模型发现不会载入预训练权重，原因是因为yolov8 model只能载入pt或者yaml，二者只能其一，我们改进模型肯定要 选择mode = yaml

(1)参考魔傀面具大佬的blibli将下面这段代码插入yolov8–>ultralytics–>yolo–>engine–>model.py–>def _new，可以载入预训练权重进行迁移学习（如果不想用预训练权重想从0开始，则注释就行）

""" 添加代码bug """
ckpt = torch.load('yolov8n.pt')
csd = ckpt['model'].float().state_dict()
csd = intersect_dicts(csd,self.model.state_dict())
self.model.load_state_dict(csd,strict = False)
print(f'Transferred {len(csd)}/{len(self.model.state_dict())} items')
""" 添加代码bug """
1
2
3
4
5
6
7

具体插入位置如下图所示

(2)yolov8 train的时候会重新载入一次模型，用yolo测试模型yaml文件一次，后面又会覆盖一次 所以这里需要修改代码 文件路径yolov8–>ultralytics–>yolo–>engine–>model.py–>def train

        self.trainer = TASK_MAP[self.task][1](overrides=overrides)
        # if not overrides.get('resume'):  # manually set model only if not resuming
        #     self.trainer.model = self.trainer.get_model(weights=self.model if self.ckpt else None, cfg=self.model.yaml)
        #     self.model = self.trainer.model

        self.trainer.model = self.model
1
2
3
4
5
6

插入具体位置如下图所示

6、然后在主目录下新建推理文件v8_predict.py，将best.pt移动到主目录下 新建v8_images文件夹存放你需要推理的图片

from ultralytics import YOLO

if __name__ == '__main__':
    # Load a model
    model = YOLO("best.pt")

    # Predict with the model
    results = model(source="v8_images",save=True,device= "cuda:0")  # predict on an image
1
2
3
4
5
6
7
8

7、2023.10.6更新，由于作者yolov8框架的更新，现在载入预训练模型可以使用作者的api接口 model.load(‘yolov8n.pt’)来使用，只要权重路径是对的，也不用使用更改源码的方式了。

from ultralytics import YOLO


if __name__ == '__main__':
    # 直接使用预训练模型创建模型.
    # model = YOLO('yolov8n.pt')
    # model.train(**{'cfg':'ultralytics/cfg/exp1.yaml', 'data':'dataset/data.yaml'})
    
    # 使用yaml配置文件来创建模型,并导入预训练权重.
    model = YOLO('ultralytics/cfg/models/v8/yolov8n.yaml')
    model.load('yolov8n.pt')
    model.train(**{'cfg':'ultralytics/cfg/PCB.yaml', 'data':'datasets/VOCPCB.yaml'})
    
    # 模型验证
    # model = YOLO('runs/detect/train11/weights/best.pt')
    # model.val(**{'cfg':'ultralytics/cfg/PCB.yaml', 'data':'datasets/VOCPCB.yaml'})
    
    # 模型推理
    # model = YOLO('runs/detect/yolov8n_exp/best.pt')
    # model.predict(source='dataset/images/test', **{'save':True})

    # 模型导出
    # model = YOLO("Weight/yolov8n.pt")  # load an official model
    # model.export(format="onnx")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25