yolov8实例分割训练自己的数据集并加载预训练模型_yolov8预训练模型

一. 引言

最近想运用实例分割算法，刚好yolov8发布了实例分割，打算使用yolov8算法跑，但是加载预训练 模型时候，发现不能，于是对代码进行了简单的修改，并成功加载预训练模型，在此记录一下过程。

二. 环境配置

Ubuntu = 18.04
cuda = 11.7
pytorch = 1.11.0
torchvision = 0.12.0
详细信息参考官方信息 requirements.txt
1
2
3
4
5

下载代码

git clone https://github.com/ultralytics/ultralytics
1

三. 数据准备

创建一个数据集目录叫mydata_seg, 按照下图目录形式存放

注意标签是： 类别 x坐标 y坐标 …（中间有空格，并且坐标除以了原图图片大小，归一化后的）这是实例分割的数据要求。

四. 修改配置文件

修改数据集，路径，类别名称
根据上面数据目录路径来修改。
vim ./ultralytics/datasets/coco-seg.yaml
path: ..mydata_seg  
train: images/train
val: images/val

names: #自己训练的名称对应上
  0: person
  1: bicycle
  2: car
  3: motorcycle
1
2
3
4
5
6
7
8
9
10
11
12

重要部分，重要部分，重要部分

修改默认配置文件，这个很重要，并在次中，添加了加载预训练的配置参数。
vim ./ultralytics/yolo/cfg/default.yaml
task: segment  # inference task, i.e. detect, segment, classify
mode: train  # YOLO mode, i.e. train, val, predict, export
model: ./ultralytics/models/v8/seg/yolov8x-seg.yaml #这个参数与官方有点不一直，只能放网络结构配置文件，不能放网络权重模型.pt
data: ./ultralytics/datasets/coco-seg.yaml #path to data file, i.e  coco-seg.yaml
weights: ./yolov8x-seg.pt #官方下载的预训练模型本地地址， 这是我们添加的参数
freenze: 8  #需要冻结的前几层的，这是我们自己添加的参数
剩下参数和官方一致根据自己的需要修改参数
epochs: 100  # number of epochs to train for
patience: 50  # epochs to wait for no observable improvement for early stopping of training
batch: 16  # number of images per batch (-1 for AutoBatch)
imgsz: 640  # size of input images as integer or w,h
...
1
2
3
4
5
6
7
8
9
10
11
12
13
14

修改网络配置文件
vim ./ultralytics/models/v8/seg/yolov8x-seg.yaml
nc: 4 #自己训练的内别总个数
1
2
3

重要部分，重要部分，重要部分

```bash
修改train.py的代码，新增加载预训练模型
vim ./ultralytics/yolo/v8/segment/train.py
导入新的包
import sys
sys.path.append("/home/bananl/ultralytics-main")
from ultralytics.yolo.utils import LOGGER
from ultralytics.yolo.utils.torch_utils import intersect_dicts
from ultralytics.yolo.utils.checks import check_yaml
from ultralytics.yolo.utils.__init__ import yaml_load

class SegmentationTrainer(v8.detect.DetectionTrainer):

    def __init__(self, cfg=DEFAULT_CFG, overrides=None):
        if overrides is None:
            overrides = {}
        overrides['task'] = 'segment'
        super().__init__(cfg, overrides)
        self.weights = overrides['weights']
        self.freenze = overrides['freenze']

    def get_model(self, cfg=None, weights=None, verbose=True):
        if self.weights:
        	ckpt = torch.load(self.weights, map_location='cpu')
        	model = SegmentationModel(cfg or ckpt['model'], ch=3, nc=self.data['nc'], verbose=verbose and RANK == -1)
        	csd = ckpt['model'].float().state_dict()
        	csd = intersect_dict(csd, model.state_dict())
        	model.load_state_dict(csd, strict=False)
        	LOGGER.info(f"Transferred {len(csd)} / {len(model.state_dict())} items from {self.weights}")
        	freeze = [self.freenze]
        	freeze = [f"model.{x}." for x in (freeze if len(freeze) > 1 else range(freeze[0]))]
        	for k, v in model.named_parameters():
        		v.requires_grad = True
        		if any(x in k for x in freeze):
        			LOGGER.info(f"freezing {k}")
        			v.requires_grad = False
        else:
          model = SegmentationModel(cfg, ch=3, nc=self.data['nc'], verbose=verbose and RANK == -1)
        return model
其余一致
...........

def train(cfg=DEFAULT_CFG, use_python=False):
    model = cfg.model or 'yolov8n-seg.pt'
    data = cfg.data or 'coco128-seg.yaml'  # or yolo.ClassificationDataset("mnist")
    freenze = cfg.freenze
    weights = cfg.weights
    device = cfg.device if cfg.device is not None else ''
    

    args = dict(model=model, data=data, device=device, weights=weights, freenze=freenze)
    if use_python:
        from ultralytics import YOLO
        YOLO(model).train(**args)
    else:
        trainer = SegmentationTrainer(overrides=args)
        trainer.train()

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
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58

开始训练

python train.py 
1

可以看出，的确成功加载了预训练模型。坑填完。