当前位置:   article > 正文

​【实战精选】MobileNetV3与YOLOv5结合:室内场景高效分割系统_yolov5 mobilenetv3

yolov5 mobilenetv3

1.研究背景与意义

近年来,随着计算机视觉技术的快速发展,室内场景分割成为了一个备受关注的研究领域。室内场景分割是指将室内环境中的不同物体和区域进行准确的分割和识别,以实现对室内场景的理解和分析。它在许多领域具有广泛的应用,如智能家居、安防监控、机器人导航等。

然而,由于室内场景的复杂性和多样性,室内场景分割面临着许多挑战。首先,室内场景中存在大量的物体和区域,它们的形状、颜色和纹理各不相同,导致分割任务的难度较大。其次,室内场景中常常存在光照不均匀、遮挡和噪声等问题,这些因素会影响分割算法的准确性和鲁棒性。此外,传统的室内场景分割方法通常需要大量的计算资源和时间,限制了其在实时应用中的应用。

为了解决上述问题,研究者们提出了许多基于深度学习的室内场景分割方法。其中,基于改进MobilenetV3骨干网络的YOLOv5是一种较为先进的方法。MobilenetV3是一种轻量级的卷积神经网络,具有较低的参数量和计算复杂度,适合在资源受限的设备上进行实时场景分割。而YOLOv5是一种快速而准确的目标检测算法,可以有效地定位和识别室内场景中的物体和区域。

本研究的主要目的是基于改进MobilenetV3骨干网络的YOLOv5,设计和实现一个高效准确的室内场景分割系统。具体来说,我们将通过以下几个方面的工作来达到目标:

首先,我们将对MobilenetV3骨干网络进行改进,以提高其在室内场景分割任务中的性能。我们将引入更多的特征层和注意力机制,以捕捉更丰富的语义信息和上下文关系。

其次,我们将结合YOLOv5的目标检测能力,将其应用于室内场景分割任务中。通过将目标检测和场景分割相结合,可以提高分割算法的准确性和鲁棒性。

最后,我们将对系统进行优化和加速,以实现实时室内场景分割。我们将采用模型压缩和量化技术,减少模型的参数量和计算复杂度,同时利用硬件加速器和并行计算等方法,提高系统的运行效率。

通过以上工作,我们期望能够设计和实现一个高效准确的室内场景分割系统,为室内环境的理解和分析提供有力的支持。该系统可以应用于智能家居、安防监控、机器人导航等领域,提高室内场景的感知和交互能力,为用户提供更好的体验和服务。同时,本研究还可以为深度学习在室内场景分割领域的应用提供一种新的思路和方法,对相关研究具有一定的参考价值。

2.图片演示

在这里插入图片描述
在这里插入图片描述
在这里插入图片描述

3.视频演示

基于改进mobilenetv3骨干网络YOLOv5的室内场景分割系统_哔哩哔哩_bilibili

4.数据集的采集&标注和整理

图片的收集

首先,我们需要收集所需的图片。这可以通过不同的方式来实现,例如使用现有的公开数据集IndoorDatasets
在这里插入图片描述

eiseg是一个图形化的图像注释工具,支持COCO和YOLO格式。以下是使用eiseg将图片标注为COCO格式的步骤:

(1)下载并安装eiseg。
(2)打开eiseg并选择“Open Dir”来选择你的图片目录。
(3)为你的目标对象设置标签名称。
(4)在图片上绘制矩形框,选择对应的标签。
(5)保存标注信息,这将在图片目录下生成一个与图片同名的JSON文件。
(6)重复此过程,直到所有的图片都标注完毕。

由于YOLO使用的是txt格式的标注,我们需要将VOC格式转换为YOLO格式。可以使用各种转换工具或脚本来实现。
在这里插入图片描述

下面是一个简单的方法是使用Python脚本,该脚本读取XML文件,然后将其转换为YOLO所需的txt格式。

import contextlib
import json

import cv2
import pandas as pd
from PIL import Image
from collections import defaultdict

from utils import *


# Convert INFOLKS JSON file into YOLO-format labels ----------------------------
def convert_infolks_json(name, files, img_path):
    # Create folders
    path = make_dirs()

    # Import json
    data = []
    for file in glob.glob(files):
        with open(file) as f:
            jdata = json.load(f)
            jdata['json_file'] = file
            data.append(jdata)

    # Write images and shapes
    name = path + os.sep + name
    file_id, file_name, wh, cat = [], [], [], []
    for x in tqdm(data, desc='Files and Shapes'):
        f = glob.glob(img_path + Path(x['json_file']).stem + '.*')[0]
        file_name.append(f)
        wh.append(exif_size(Image.open(f)))  # (width, height)
        cat.extend(a['classTitle'].lower() for a in x['output']['objects'])  # categories

        # filename
        with open(name + '.txt', 'a') as file:
            file.write('%s\n' % f)

    # Write *.names file
    names = sorted(np.unique(cat))
    # names.pop(names.index('Missing product'))  # remove
    with open(name + '.names', 'a') as file:
        [file.write('%s\n' % a) for a in names]

    # Write labels file
    for i, x in enumerate(tqdm(data, desc='Annotations')):
        label_name = Path(file_name[i]).stem + '.txt'

        with open(path + '/labels/' + label_name, 'a') as file:
            for a in x['output']['objects']:
                # if a['classTitle'] == 'Missing product':
                #    continue  # skip

                category_id = names.index(a['classTitle'].lower())

                # The INFOLKS bounding box format is [x-min, y-min, x-max, y-max]
                box = np.array(a['points']['exterior'], dtype=np.float32).ravel()
                box[[0, 2]] /= wh[i][0]  # normalize x by width
                box[[1, 3]] /= wh[i][1]  # normalize y by height
                box = [box[[0, 2]].mean(), box[[1, 3]].mean(), box[2] - box[0], box[3] - box[1]]  # xywh
                if (box[2] > 0.) and (box[3] > 0.):  # if w > 0 and h > 0
                    file.write('%g %.6f %.6f %.6f %.6f\n' % (category_id, *box))

    # Split data into train, test, and validate files
    split_files(name, file_name)
    write_data_data(name + '.data', nc=len(names))
    print(f'Done. Output saved to {os.getcwd() + os.sep + path}')


# Convert vott JSON file into YOLO-format labels -------------------------------
def convert_vott_json(name, files, img_path):
    # Create folders
    path = make_dirs()
    name = path + os.sep + name

    # Import json
    data = []
    for file in glob.glob(files):
        with open(file) as f:
            jdata = json.load(f)
            jdata['json_file'] = file
            data.append(jdata)

    # Get all categories
    file_name, wh, cat = [], [], []
    for i, x in enumerate(tqdm(data, desc='Files and Shapes')):
        with contextlib.suppress(Exception):
            cat.extend(a['tags'][0] for a in x['regions'])  # categories

    # Write *.names file
    names = sorted(pd.unique(cat))
    with open(name + '.names', 'a') as file:
        [file.write('%s\n' % a) for a in names]

    # Write labels file
    n1, n2 = 0, 0
    missing_images = []
    for i, x in enumerate(tqdm(data, desc='Annotations')):

        f = glob.glob(img_path + x['asset']['name'] + '.jpg')
        if len(f):
            f = f[0]
            file_name.append(f)
            wh = exif_size(Image.open(f))  # (width, height)

            n1 += 1
            if (len(f) > 0) and (wh[0] > 0) and (wh[1] > 0):
                n2 += 1

                # append filename to list
                with open(name + '.txt', 'a') as file:
                    file.write('%s\n' % f)

                # write labelsfile
                label_name = Path(f).stem + '.txt'
                with open(path + '/labels/' + label_name, 'a') as file:
                    for a in x['regions']:
                        category_id = names.index(a['tags'][0])

                        # The INFOLKS bounding box format is [x-min, y-min, x-max, y-max]
                        box = a['boundingBox']
                        box = np.array([box['left'], box['top'], box['width'], box['height']]).ravel()
                        box[[0, 2]] /= wh[0]  # normalize x by width
                        box[[1, 3]] /= wh[1]  # normalize y by height
                        box = [box[0] + box[2] / 2, box[1] + box[3] / 2, box[2], box[3]]  # xywh

                        if (box[2] > 0.) and (box[3] > 0.):  # if w > 0 and h > 0
                            file.write('%g %.6f %.6f %.6f %.6f\n' % (category_id, *box))
        else:
            missing_images.append(x['asset']['name'])

    print('Attempted %g json imports, found %g images, imported %g annotations successfully' % (i, n1, n2))
    if len(missing_images):
        print('WARNING, missing images:', missing_images)

    # Split data into train, test, and validate files
    split_files(name, file_name)
    print(f'Done. Output saved to {os.getcwd() + os.sep + path}')


# Convert ath JSON file into YOLO-format labels --------------------------------
def convert_ath_json(json_dir):  # dir contains json annotations and images
    # Create folders
    dir = make_dirs()  # output directory

    jsons = []
    for dirpath, dirnames, filenames in os.walk(json_dir):
        jsons.extend(
            os.path.join(dirpath, filename)
            for filename in [
                f for f in filenames if f.lower().endswith('.json')
            ]
        )

    # Import json
    n1, n2, n3 = 0, 0, 0
    missing_images, file_name = [], []
    for json_file in sorted(jsons):
        with open(json_file) as f:
            data = json.load(f)

        # # Get classes
        # try:
        #     classes = list(data['_via_attributes']['region']['class']['options'].values())  # classes
        # except:
        #     classes = list(data['_via_attributes']['region']['Class']['options'].values())  # classes

        # # Write *.names file
        # names = pd.unique(classes)  # preserves sort order
        # with open(dir + 'data.names', 'w') as f:
        #     [f.write('%s\n' % a) for a in names]

        # Write labels file
        for x in tqdm(data['_via_img_metadata'].values(), desc=f'Processing {json_file}'):
            image_file = str(Path(json_file).parent / x['filename'])
            f = glob.glob(image_file)  # image file
            if len(f):
                f = f[0]
                file_name.append(f)
                wh = exif_size(Image.open(f))  # (width, height)

                n1 += 1  # all images
                if len(f) > 0 and wh[0] > 0 and wh[1] > 0:
                    label_file = dir + 'labels/' + Path(f).stem + '.txt'

                    nlabels = 0
                    try:
                        with open(label_file, 'a') as file:  # write labelsfile
                            # try:
                            #     category_id = int(a['region_attributes']['class'])
                            # except:
                            #     category_id = int(a['region_attributes']['Class'])
                            category_id = 0  # single-class

                            for a in x['regions']:
                                # bounding box format is [x-min, y-min, x-max, y-max]
                                box = a['shape_attributes']
                                box = np.array([box['x'], box['y'], box['width'], box['height']],
                                               dtype=np.float32).ravel()
                                box[[0, 2]] /= wh[0]  # normalize x by width
                                box[[1, 3]] /= wh[1]  # normalize y by height
                                box = [box[0] + box[2] / 2, box[1] + box[3] / 2, box[2],
                                       box[3]]  # xywh (left-top to center x-y)

                                if box[2] > 0. and box[3] > 0.:  # if w > 0 and h > 0
                                    file.write('%g %.6f %.6f %.6f %.6f\n' % (category_id, *box))
                                    n3 += 1
                                    nlabels += 1

                        if nlabels == 0:  # remove non-labelled images from dataset
                            os.system(f'rm {label_file}')
                            # print('no labels for %s' % f)
                            continue  # next file

                        # write image
                        img_size = 4096  # resize to maximum
                        img = cv2.imread(f)  # BGR
                        assert img is not None, 'Image Not Found ' + f
                        r = img_size / max(img.shape)  # size ratio
                        if r < 1:  # downsize if necessary
                            h, w, _ = img.shape
                            img = cv2.resize(img, (int(w * r), int(h * r)), interpolation=cv2.INTER_AREA)

                        ifile = dir + 'images/' + Path(f).name
                        if cv2.imwrite(ifile, img):  # if success append image to list
                            with open(dir + 'data.txt', 'a') as file:
                                file.write('%s\n' % ifile)
                            n2 += 1  # correct images

                    except Exception:
                        os.system(f'rm {label_file}')
                        print(f'problem with {f}')

            else:
                missing_images.append(image_file)

    nm = len(missing_images)  # number missing
    print('\nFound %g JSONs with %g labels over %g images. Found %g images, labelled %g images successfully' %
          (len(jsons), n3, n1, n1 - nm, n2))
    if len(missing_images):
        print('WARNING, missing images:', missing_images)

    # Write *.names file
    names = ['knife']  # preserves sort order
    with open(dir + 'data.names', 'w') as f:
        [f.write('%s\n' % a) for a in names]

    # Split data into train, test, and validate files
    split_rows_simple(dir + 'data.txt')
    write_data_data(dir + 'data.data', nc=1)
    print(f'Done. Output saved to {Path(dir).absolute()}')


def convert_coco_json(json_dir='../coco/annotations/', use_segments=False, cls91to80=False):
    save_dir = make_dirs()  # output directory
    coco80 = coco91_to_coco80_class()

    # Import json
    for json_file in sorted(Path(json_dir).resolve().glob('*.json')):
        fn = Path(save_dir) / 'labels' / json_file.stem.replace('instances_', '')  # folder name
        fn.mkdir()
        with open(json_file) as f:
            data = json.load(f)

        # Create image dict
        images = {'%g' % x['id']: x for x in data['images']}
        # Create image-annotations dict
        imgToAnns = defaultdict(list)
        for ann in data['annotations']:
            imgToAnns[ann['image_id']].append(ann)

        # Write labels file
        for img_id, anns in tqdm(imgToAnns.items(), desc=f'Annotations {json_file}'):
            img = images['%g' % img_id]
            h, w, f = img['height'], img['width'], img['file_name']

            bboxes = []
            segments = []
            for ann in anns:
                if ann['iscrowd']:
                    continue
                # The COCO box format is [top left x, top left y, width, height]
                box = np.array(ann['bbox'], dtype=np.float64)
                box[:2] += box[2:] / 2  # xy top-left corner to center
                box[[0, 2]] /= w  # normalize x
                box[[1, 3]] /= h  # normalize y
                if box[2] <= 0 or box[3] <= 0:  # if w <= 0 and h <= 0
                    continue

                cls = coco80[ann['category_id'] - 1] if cls91to80 else ann['category_id'] - 1  # class
                box = [cls] + box.tolist()
                if box not in bboxes:
                    bboxes.append(box)
                # Segments
                if use_segments:
                    if len(ann['segmentation']) > 1:
                        s = merge_multi_segment(ann['segmentation'])
                        s = (np.concatenate(s, axis=0) / np.array([w, h])).reshape(-1).tolist()
                    else:
                        s = [j for i in ann['segmentation'] for j in i]  # all segments concatenated
                        s = (np.array(s).reshape(-1, 2) / np.array([w, h])).reshape(-1).tolist()
                    s = [cls] + s
                    if s not in segments:
                        segments.append(s)

            # Write
            with open((fn / f).with_suffix('.txt'), 'a') as file:
                for i in range(len(bboxes)):
                    line = *(segments[i] if use_segments else bboxes[i]),  # cls, box or segments
                    file.write(('%g ' * len(line)).rstrip() % line + '\n')


def min_index(arr1, arr2):
    """Find a pair of indexes with the shortest distance. 
    Args:
        arr1: (N, 2).
        arr2: (M, 2).
    Return:
        a pair of indexes(tuple).
    """
    dis = ((arr1[:, None, :] - arr2[None, :, :]) ** 2).sum(-1)
    return np.unravel_index(np.argmin(dis, axis=None), dis.shape)


def merge_multi_segment(segments):
    """Merge multi segments to one list.
    Find the coordinates with min distance between each segment,
    then connect these coordinates with one thin line to merge all 
    segments into one.

    Args:
        segments(List(List)): original segmentations in coco's json file.
            like [segmentation1, segmentation2,...], 
            each segmentation is a list of coordinates.
    """
    s = []
    segments = [np.array(i).reshape(-1, 2) for i in segments]
    idx_list = [[] for _ in range(len(segments))]

    # record the indexes with min distance between each segment
    for i in range(1, len(segments)):
        idx1, idx2 = min_index(segments[i - 1], segments[i])
        idx_list[i - 1].append(idx1)
        idx_list[i].append(idx2)

    # use two round to connect all the segments
    for k in range(2):
        # forward connection
        if k == 0:
            for i, idx in enumerate(idx_list):
                # middle segments have two indexes
                # reverse the index of middle segments
                if len(idx) == 2 and idx[0] > idx[1]:
                    idx = idx[::-1]
                    segments[i] = segments[i][::-1, :]

                segments[i] = np.roll(segments[i], -idx[0], axis=0)
                segments[i] = np.concatenate([segments[i], segments[i][:1]])
                # deal with the first segment and the last one
                if i in [0, len(idx_list) - 1]:
                    s.append(segments[i])
                else:
                    idx = [0, idx[1] - idx[0]]
                    s.append(segments[i][idx[0]:idx[1] + 1])

        else:
            for i in range(len(idx_list) - 1, -1, -1):
                if i not in [0, len(idx_list) - 1]:
                    idx = idx_list[i]
                    nidx = abs(idx[1] - idx[0])
                    s.append(segments[i][nidx:])
    return s


def delete_dsstore(path='../datasets'):
    # Delete apple .DS_store files
    from pathlib import Path
    files = list(Path(path).rglob('.DS_store'))
    print(files)
    for f in files:
        f.unlink()


if __name__ == '__main__':
    source = 'COCO'

    if source == 'COCO':
        convert_coco_json('./annotations',  # directory with *.json
                          use_segments=True,
                          cls91to80=True)

    elif source == 'infolks':  # Infolks https://infolks.info/
        convert_infolks_json(name='out',
                             files='../data/sm4/json/*.json',
                             img_path='../data/sm4/images/')

    elif source == 'vott':  # VoTT https://github.com/microsoft/VoTT
        convert_vott_json(name='data',
                          files='../../Downloads/athena_day/20190715/*.json',
                          img_path='../../Downloads/athena_day/20190715/')  # images folder

    elif source == 'ath':  # ath format
        convert_ath_json(json_dir='../../Downloads/athena/')  # images folder

    # zip results
    # os.system('zip -r ../coco.zip ../coco')


  • 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
  • 59
  • 60
  • 61
  • 62
  • 63
  • 64
  • 65
  • 66
  • 67
  • 68
  • 69
  • 70
  • 71
  • 72
  • 73
  • 74
  • 75
  • 76
  • 77
  • 78
  • 79
  • 80
  • 81
  • 82
  • 83
  • 84
  • 85
  • 86
  • 87
  • 88
  • 89
  • 90
  • 91
  • 92
  • 93
  • 94
  • 95
  • 96
  • 97
  • 98
  • 99
  • 100
  • 101
  • 102
  • 103
  • 104
  • 105
  • 106
  • 107
  • 108
  • 109
  • 110
  • 111
  • 112
  • 113
  • 114
  • 115
  • 116
  • 117
  • 118
  • 119
  • 120
  • 121
  • 122
  • 123
  • 124
  • 125
  • 126
  • 127
  • 128
  • 129
  • 130
  • 131
  • 132
  • 133
  • 134
  • 135
  • 136
  • 137
  • 138
  • 139
  • 140
  • 141
  • 142
  • 143
  • 144
  • 145
  • 146
  • 147
  • 148
  • 149
  • 150
  • 151
  • 152
  • 153
  • 154
  • 155
  • 156
  • 157
  • 158
  • 159
  • 160
  • 161
  • 162
  • 163
  • 164
  • 165
  • 166
  • 167
  • 168
  • 169
  • 170
  • 171
  • 172
  • 173
  • 174
  • 175
  • 176
  • 177
  • 178
  • 179
  • 180
  • 181
  • 182
  • 183
  • 184
  • 185
  • 186
  • 187
  • 188
  • 189
  • 190
  • 191
  • 192
  • 193
  • 194
  • 195
  • 196
  • 197
  • 198
  • 199
  • 200
  • 201
  • 202
  • 203
  • 204
  • 205
  • 206
  • 207
  • 208
  • 209
  • 210
  • 211
  • 212
  • 213
  • 214
  • 215
  • 216
  • 217
  • 218
  • 219
  • 220
  • 221
  • 222
  • 223
  • 224
  • 225
  • 226
  • 227
  • 228
  • 229
  • 230
  • 231
  • 232
  • 233
  • 234
  • 235
  • 236
  • 237
  • 238
  • 239
  • 240
  • 241
  • 242
  • 243
  • 244
  • 245
  • 246
  • 247
  • 248
  • 249
  • 250
  • 251
  • 252
  • 253
  • 254
  • 255
  • 256
  • 257
  • 258
  • 259
  • 260
  • 261
  • 262
  • 263
  • 264
  • 265
  • 266
  • 267
  • 268
  • 269
  • 270
  • 271
  • 272
  • 273
  • 274
  • 275
  • 276
  • 277
  • 278
  • 279
  • 280
  • 281
  • 282
  • 283
  • 284
  • 285
  • 286
  • 287
  • 288
  • 289
  • 290
  • 291
  • 292
  • 293
  • 294
  • 295
  • 296
  • 297
  • 298
  • 299
  • 300
  • 301
  • 302
  • 303
  • 304
  • 305
  • 306
  • 307
  • 308
  • 309
  • 310
  • 311
  • 312
  • 313
  • 314
  • 315
  • 316
  • 317
  • 318
  • 319
  • 320
  • 321
  • 322
  • 323
  • 324
  • 325
  • 326
  • 327
  • 328
  • 329
  • 330
  • 331
  • 332
  • 333
  • 334
  • 335
  • 336
  • 337
  • 338
  • 339
  • 340
  • 341
  • 342
  • 343
  • 344
  • 345
  • 346
  • 347
  • 348
  • 349
  • 350
  • 351
  • 352
  • 353
  • 354
  • 355
  • 356
  • 357
  • 358
  • 359
  • 360
  • 361
  • 362
  • 363
  • 364
  • 365
  • 366
  • 367
  • 368
  • 369
  • 370
  • 371
  • 372
  • 373
  • 374
  • 375
  • 376
  • 377
  • 378
  • 379
  • 380
  • 381
  • 382
  • 383
  • 384
  • 385
  • 386
  • 387
  • 388
  • 389
  • 390
  • 391
  • 392
  • 393
  • 394
  • 395
  • 396
  • 397
  • 398
  • 399
  • 400
  • 401
  • 402
  • 403
  • 404
  • 405
  • 406
  • 407
整理数据文件夹结构

我们需要将数据集整理为以下结构:

-----datasets
	-----coco128-seg
	   |-----images
	   |   |-----train
	   |   |-----valid
	   |   |-----test
	   |
	   |-----labels
	   |   |-----train
	   |   |-----valid
	   |   |-----test
	   |

  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
模型训练
 Epoch   gpu_mem       box       obj       cls    labels  img_size
 1/200     20.8G   0.01576   0.01955  0.007536        22      1280: 100%|██████████| 849/849 [14:42<00:00,  1.04s/it]
           Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|██████████| 213/213 [01:14<00:00,  2.87it/s]
             all       3395      17314      0.994      0.957      0.0957      0.0843

 Epoch   gpu_mem       box       obj       cls    labels  img_size
 2/200     20.8G   0.01578   0.01923  0.007006        22      1280: 100%|██████████| 849/849 [14:44<00:00,  1.04s/it]
           Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|██████████| 213/213 [01:12<00:00,  2.95it/s]
             all       3395      17314      0.996      0.956      0.0957      0.0845

 Epoch   gpu_mem       box       obj       cls    labels  img_size
 3/200     20.8G   0.01561    0.0191  0.006895        27      1280: 100%|██████████| 849/849 [10:56<00:00,  1.29it/s]
           Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|███████   | 187/213 [00:52<00:00,  4.04it/s]
             all       3395      17314      0.996      0.957      0.0957      0.0845
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6
  • 7
  • 8
  • 9
  • 10
  • 11
  • 12
  • 13
  • 14

5.核心代码讲解

5.1 export.py


def export_formats():
    # YOLOv5 export formats
    x = [
        ['PyTorch', '-', '.pt', True, True],
        ['TorchScript', 'torchscript', '.torchscript', True, True],
        ['ONNX', 'onnx', '.onnx', True, True],
        ['OpenVINO', 'openvino', '_openvino_model', True, False],
        ['TensorRT', 'engine', '.engine', False, True],
        ['CoreML', 'coreml', '.mlmodel', True, False],
        ['TensorFlow SavedModel', 'saved_model', '_saved_model', True, True],
        ['TensorFlow GraphDef', 'pb', '.pb', True, True],
        ['TensorFlow Lite', 'tflite', '.tflite', True, False],
        ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False, False],
        ['TensorFlow.js', 'tfjs', '_web_model', False, False],
        ['PaddlePaddle', 'paddle', '_paddle_model', True, True],]
    return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'CPU', 'GPU'])


def try_export(inner_func):
    # YOLOv5 export decorator, i..e @try_export
    inner_args = get_default_args(inner_func)

    def outer_func(*args, **kwargs):
        prefix = inner_args['prefix']
        try:
            with Profile() as dt:
                f, model = inner_func(*args, **kwargs)
            LOGGER.info(f'{prefix} export success ✅ {dt.t:.1f}s, saved as {f} ({file_size(f):.1f} MB)')
            return f, model
        except Exception as e:
            LOGGER.info(f'{prefix} export failure ❌ {dt.t:.1f}s: {e}')
            return None, None

    return outer_func


@try_export
def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')):
    # YOLOv5 TorchScript model export
    LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...')
    f = file.with_suffix('.torchscript')

    ts = torch.jit.trace(model, im, strict=False)
    d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names}
    extra_files = {'config.txt': json.dumps(d)}  # torch._C.ExtraFilesMap()
    if optimize:  # https://pytorch.org/tutorials/recipes/mobile_interpreter.html
        optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files)
    else:
        ts.save(str(f), _extra_files=extra_files)
    return f, None


@try_export
def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX:')):
    # YOLOv5 ONNX export
    check_requirements('onnx>=1.12.0')
    import onnx

    LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
    f = file.with_suffix('.onnx')

    output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output0']
    if dynamic:
        dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}}  # shape(1,3,640,640)
        if isinstance(model, SegmentationModel):
            dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
            dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'}  # shape(1,32,160,160)
        elif isinstance(model, DetectionModel):
            dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)

    torch.onnx.export(
        model.cpu() if dynamic else model,  # --dynamic only compatible with cpu
        im.cpu() if dynamic else im,
        f,
        verbose=False,
        opset_version=opset,
        do_constant_folding=True,  # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
        input_names=['images'],
        output_names=output_names,
        dynamic_axes=dynamic or None)

    # Checks
    model_onnx = onnx.load(f)  # load onnx model
    onnx.checker.check_model(model_onnx)  # check onnx model

    # Metadata
    d = {'stride': int(max(model.stride)), 'names': model.names}
    for k, v in d.items():
        meta = model_onnx.metadata_props.add()
        meta.key, meta.value = k, str(v)
    onnx.save(model_onnx, f)

    # Simplify
    if simplify:
        try:
            cuda = torch.cuda.is_available()
            check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1'))
            import onnxsim

            LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
            model_simp, check = onnxsim.simplify(f, check=True)
            assert check, 'assert check failed'
            onnx.save(model_simp, f)
        except Exception as e:
            LOGGER.info(f'{prefix} simplifier failure {e}')
    return f, None
  • 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
  • 59
  • 60
  • 61
  • 62
  • 63
  • 64
  • 65
  • 66
  • 67
  • 68
  • 69
  • 70
  • 71
  • 72
  • 73
  • 74
  • 75
  • 76
  • 77
  • 78
  • 79
  • 80
  • 81
  • 82
  • 83
  • 84
  • 85
  • 86
  • 87
  • 88
  • 89
  • 90
  • 91
  • 92
  • 93
  • 94
  • 95
  • 96
  • 97
  • 98
  • 99
  • 100
  • 101
  • 102
  • 103
  • 104
  • 105
  • 106
  • 107
  • 108

export.py是一个用于将YOLOv5 PyTorch模型导出为其他格式的程序文件。它支持导出的格式包括PyTorch、TorchScript、ONNX、OpenVINO、TensorRT、CoreML、TensorFlow SavedModel、TensorFlow GraphDef、TensorFlow Lite、TensorFlow Edge TPU、TensorFlow.js和PaddlePaddle。该文件还包含了一些导出所需的依赖库和使用示例。

程序文件首先定义了一些常量和全局变量,然后定义了一些辅助函数和装饰器。其中,辅助函数export_formats()返回了一个DataFrame,包含了YOLOv5支持的导出格式的相关信息。装饰器try_export()用于捕获导出过程中的异常,并在导出成功或失败时打印相应的信息。

接下来,程序文件定义了一些具体的导出函数,如export_torchscript()用于导出TorchScript模型,export_onnx()用于导出ONNX模型。这些函数使用了PyTorch和相关的库来实现导出功能。

最后,程序文件定义了一个主函数,用于解析命令行参数并执行相应的导出操作。该函数首先加载YOLOv5模型,然后根据命令行参数选择要导出的格式,并调用相应的导出函数进行导出。导出成功后,打印导出的文件路径和大小。

整个程序文件的功能是将YOLOv5 PyTorch模型导出为其他格式,以便在不同的平台和框架上进行推理和部署。

5.2 mobilenetv3.py


class hswish(nn.Module):
    def forward(self, x):
        out = x * F.relu6(x + 3, inplace=True) / 6
        return out

class hsigmoid(nn.Module):
    def forward(self, x):
        out = F.relu6(x + 3, inplace=True) / 6
        return out

class SeModule(nn.Module):
    def __init__(self, in_size, reduction=4):
        super(SeModule, self).__init__()
        expand_size =  max(in_size // reduction, 8)
        self.se = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_size, expand_size, kernel_size=1, bias=False),
            nn.BatchNorm2d(expand_size),
            nn.ReLU(inplace=True),
            nn.Conv2d(expand_size, in_size, kernel_size=1, bias=False),
            nn.Hardsigmoid()
        )

    def forward(self, x):
        return x * self.se(x)

class Block(nn.Module):
    '''expand + depthwise + pointwise'''
    def __init__(self, kernel_size, in_size, expand_size, out_size, act, se, stride):
        super(Block, self).__init__()
        self.stride = stride

        self.conv1 = nn.Conv2d(in_size, expand_size, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(expand_size)
        self.act1 = act(inplace=True)

        self.conv2 = nn.Conv2d(expand_size, expand_size, kernel_size=kernel_size, stride=stride, padding=kernel_size//2, groups=expand_size, bias=False)
        self.bn2 = nn.BatchNorm2d(expand_size)
        self.act2 = act(inplace=True)
        self.se = SeModule(expand_size) if se else nn.Identity()

        self.conv3 = nn.Conv2d(expand_size, out_size, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(out_size)
        self.act3 = act(inplace=True)

        self.skip = None
        if stride == 1 and in_size != out_size:
            self.skip = nn.Sequential(
                nn.Conv2d(in_size, out_size, kernel_size=1, bias=False),
                nn.BatchNorm2d(out_size)
            )

        if stride == 2 and in_size != out_size:
            self.skip = nn.Sequential(
                nn.Conv2d(in_channels=in_size, out_channels=in_size, kernel_size=3, groups=in_size, stride=2, padding=1, bias=False),
                nn.BatchNorm2d(in_size),
                nn.Conv2d(in_size, out_size, kernel_size=1, bias=True),
                nn.BatchNorm2d(out_size)
            )

        if stride == 2 and in_size == out_size:
            self.skip = nn.Sequential(
                nn.Conv2d(in_channels=in_size, out_channels=out_size, kernel_size=3, groups=in_size, stride=2, padding=1, bias=False),
                nn.BatchNorm2d(out_size)
            )

    def forward(self, x):
        skip = x

        out = self.act1(self.bn1(self.conv1(x)))
        out = self.act2(self.bn2(self.conv2(out)))
        out = self.se(out)
        out = self.bn3(self.conv3(out))
        
        if self.skip is not None:
            skip = self.skip(skip)
        return self.act3(out + skip)

class MobileNetV3_Small(nn.Module):
    def __init__(self, num_classes=1000, act=nn.Hardswish):
        super(MobileNetV3_Small, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(16)
        self.hs1 = act(inplace=True)

        self.bneck = nn.Sequential(
            Block(3, 16, 16, 16, nn.ReLU, True, 2),
            Block(3, 16, 72, 24, nn.ReLU, False, 2),
            Block(3, 24, 88, 24, nn.ReLU, False, 1),
            Block(5, 24, 96, 40, act, True, 2),
            Block(5, 40, 240, 40, act, True, 1),
            Block(5, 40, 240, 40, act, True, 1),
            Block(5, 40, 120, 48, act, True, 1),
            Block(5, 48, 144, 48, act, True, 1),
            Block(5, 48, 288, 96, act, True, 2),
            Block(5, 96, 576, 96, act, True, 1),
            Block(5, 96, 576, 96, act, True, 1),
        )

        self.conv2 = nn.Conv2d(96, 576, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn2 = nn.BatchNorm2d(576)
        self.hs2 = act(inplace=True)
        self.gap = nn.AdaptiveAvgPool2d(1)

        self.linear3 = nn.Linear(576, 1280, bias=False)
        self.bn3 = nn.BatchNorm1d(1280)
        self.hs3 = act(inplace=True)
        self.drop = nn.Dropout(0.2)
        self.linear4 = nn.Linear(1280, num_classes)
        self.init_params()

    def init_params(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, x):
        out = self.hs1(self.bn1(self.conv1(x)))
        out = self.bneck(out)

        out = self.hs2(self.bn2(self.conv2(out)))
        out = self.gap(out).flatten(1)
        out = self.drop(self.hs3(self.bn3(self.linear3(out))))

        return self.linear4(out)


class MobileNetV3_Large(nn.Module):
    def __init__(self, num_classes=1000, act=nn.Hardswish):
        super(MobileNetV3_Large, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(16)
        self.hs1 = act(inplace=True)

        self.bneck = nn.Sequential(
            Block(3, 16, 16, 16, nn.ReLU, False, 1),
            Block(3, 16, 64, 24, nn.ReLU, False, 2),
            Block(3, 24, 72, 24, nn.ReLU, False, 1),
            Block(5, 24, 72, 40, nn.ReLU, True, 2),
            Block(5, 40, 120, 40, nn.ReLU, True, 1),
            Block(5, 40, 120, 40, nn.ReLU, True, 1),
            Block(3, 40, 240, 80, act, False, 2),
            Block(3, 80, 200, 80, act, False, 1),
            Block(3, 80, 184, 80, act, False, 1),
            Block(3, 80, 184, 80, act, False, 1),
            Block(3, 80, 480, 112, act, True, 1),
            Block(3, 112, 672, 112, act, True, 1),
            Block(5, 112, 672, 160, act, True, 2),
            ......
  • 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
  • 59
  • 60
  • 61
  • 62
  • 63
  • 64
  • 65
  • 66
  • 67
  • 68
  • 69
  • 70
  • 71
  • 72
  • 73
  • 74
  • 75
  • 76
  • 77
  • 78
  • 79
  • 80
  • 81
  • 82
  • 83
  • 84
  • 85
  • 86
  • 87
  • 88
  • 89
  • 90
  • 91
  • 92
  • 93
  • 94
  • 95
  • 96
  • 97
  • 98
  • 99
  • 100
  • 101
  • 102
  • 103
  • 104
  • 105
  • 106
  • 107
  • 108
  • 109
  • 110
  • 111
  • 112
  • 113
  • 114
  • 115
  • 116
  • 117
  • 118
  • 119
  • 120
  • 121
  • 122
  • 123
  • 124
  • 125
  • 126
  • 127
  • 128
  • 129
  • 130
  • 131
  • 132
  • 133
  • 134
  • 135
  • 136
  • 137
  • 138
  • 139
  • 140
  • 141
  • 142
  • 143
  • 144
  • 145
  • 146
  • 147
  • 148
  • 149
  • 150
  • 151
  • 152
  • 153
  • 154
  • 155
  • 156
  • 157
  • 158
  • 159
  • 160

这个程序文件是一个使用PyTorch实现的MobileNetV3模型。它包含了两个不同大小的MobileNetV3模型:MobileNetV3_Small和MobileNetV3_Large。

MobileNetV3_Small模型包含了一系列的Block,每个Block由一个扩展卷积层、一个深度卷积层、一个点卷积层和一个SE模块组成。SE模块用于增强特征的表示能力。最后通过全局平均池化层和全连接层得到最终的分类结果。

MobileNetV3_Large模型与MobileNetV3_Small模型类似,但是具有更多的Block和更大的通道数。

这个程序文件还定义了一些辅助的模块,如hswish和hsigmoid激活函数,以及SeModule和Block模块。

在模型的forward方法中,输入的图像经过一系列的卷积、池化和全连接操作,最后得到分类结果。

整个程序文件的目的是实现MobileNetV3模型,用于图像分类任务。

5.3 model.py


class ModelProfiler:
    def __init__(self, model_name, input_shape):
        self.model_name = model_name
        self.input_shape = input_shape
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = None

    def load_model(self):
        self.model = timm.create_model(self.model_name, pretrained=False, features_only=True)
        self.model.to(self.device)
        self.model.eval()

    def print_model_info(self):
        print(self.model.feature_info.channels())
        for feature in self.model(self.dummy_input):
            print(feature.size())

    def profile_model(self):
        flops, params = profile(self.model.to(self.device), (self.dummy_input,), verbose=False)
        flops, params = clever_format([flops * 2, params], "%.3f")
        print('Total FLOPS: %s' % (flops))
        print('Total params: %s' % (params))

    def run(self):
        self.load_model()
        self.print_model_info()
        self.profile_model()

model_name = 'mobilenetv3_small_100'
input_shape = (1, 3, 640, 640)

profiler = ModelProfiler(model_name, input_shape)
profiler.run()
  • 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

这个程序文件名为model.py,主要功能是使用timm库中的模型来计算模型的FLOPS和参数数量。

首先,程序列出了timm库中可用的所有模型。然后,它创建了一个虚拟输入张量dummy_input,并将其发送到GPU上进行计算。

接下来,程序选择了一个名为’mobilenetv3_small_100’的模型,并将其加载到设备上。然后,它将模型设置为评估模式。

程序打印了模型的特征通道数,并使用虚拟输入张量计算了模型的输出。每个特征的大小也被打印出来。

最后,程序使用thop库中的profile函数计算了模型的FLOPS和参数数量,并使用clever_format函数将其格式化为易读的形式。计算结果被打印出来。

5.4 train.py



def train(hyp,  # path/to/hyp.yaml or hyp dictionary
          opt,
          device,
          callbacks
          ):
    save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze, = \
        Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \
        opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze

    # Directories
    w = save_dir / 'weights'  # weights dir
    (w.parent if evolve else w).mkdir(parents=True, exist_ok=True)  # make dir
    last, best = w / 'last.pt', w / 'best.pt'

    # Hyperparameters
    if isinstance(hyp, str):
        with open(hyp, errors='ignore') as f:
            hyp = yaml.safe_load(f)  # load hyps dict
    LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))

    # Save run settings
    with open(save_dir / 'hyp.yaml', 'w') as f:
        yaml.safe_dump(hyp, f, sort_keys=False)
    with open(save_dir / 'opt.yaml', 'w') as f:
        yaml.safe_dump(vars(opt), f, sort_keys=False)
    data_dict = None

    # Loggers
    if RANK in [-1, 0]:
        loggers = Loggers(save_dir, weights, opt, hyp, LOGGER)  # loggers instance
        if loggers.wandb:
            data_dict = loggers.wandb.data_dict
            if resume:
                weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp

        # Register actions
        for k in methods(loggers):
            callbacks.register_action(k, callback=getattr(loggers, k))

    # Config
    plots = not evolve  # create plots
    cuda = device.type != 'cpu'
    init_seeds(1 + RANK)
    with torch_distributed_zero_first(LOCAL_RANK):
        data_dict = data_dict or check_dataset(data)  # check if None
    train_path, val_path = data_dict['train'], data_dict['val']
    nc = 1 if single_cls else int(data_dict['nc'])  # number of classes
    names = ['item'] if single_cls and len(data_dict['names']) != 1 else data_dict['names']  # class names
    assert len(names) == nc, f'{len(names)} names found for nc={nc} dataset in {data}'  # check
    is_coco = data.endswith('coco.yaml') and nc == 80  # COCO dataset

    # Model
    check_suffix(weights, '.pt')  # check weights
    pretrained = weights.endswith('.pt')
    if pretrained:
        with torch_distributed_zero_first(LOCAL_RANK):
            weights = attempt_download(weights)  # download if not found locally
        ckpt = torch.load(weights, map_location=device)  # load checkpoint
        model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
        exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else []  # exclude keys
        csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32
        csd = intersect_dicts(csd, model.state_dict(), exclude=exclude)  # intersect
        model.load_state_dict(csd, strict=False)  # load
        LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}')  # report
    else:
        model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create

    # Freeze
    freeze = [f'model.{x}.' for x in range(freeze)]  # layers to freeze
    for k, v in model.named_parameters():
        v.requires_grad = True  # train all layers
        if any(x in k for x in freeze):
            print(f'freezing {k}')
            v.requires_grad = False

    # Image size
    gs = max(int(model.stride.max()), 32)  # grid size (max stride)
    imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2)  # verify imgsz is gs-multiple

    # Batch size
    if RANK == -1 and batch_size == -1:  # single-GPU only, estimate best batch size
        batch_size = check_train_batch_size(model, imgsz)

    # Optimizer
    nbs = 64  # nominal batch size
    accumulate = max(round(nbs / batch_size), 1)  # accumulate loss before optimizing
    hyp['weight_decay'] *= batch_size * accumulate / nbs  # scale weight_decay
    LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}")

    g0, g1, g2 = [], [], []  # optimizer parameter groups
    for v in model.modules():
        if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):  # bias
            g2.append(v.bias)
        if isinstance(v, nn.BatchNorm2d):  # weight (no decay)
            g0.append(v.weight)
        elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):  # weight (with decay)
            g1.append(v.weight)

    if opt.adam:
        optimizer = Adam(g0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999))  # adjust beta1 to momentum
    else:
        optimizer = SGD(g0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)

    optimizer.add_param_group({'params': g1, 'weight_decay': hyp['weight_decay']})  # add g1 with weight_decay
    optimizer.add_param_group({'params': g2})  # add g2 (biases)
    LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__} with parameter groups "
                f"{len(g0)} weight, {len(g1)} weight (no decay), {len(g2)} bias")
    del g0, g1, g2

    # Scheduler
    if opt.linear_lr:
        ......
  • 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
  • 59
  • 60
  • 61
  • 62
  • 63
  • 64
  • 65
  • 66
  • 67
  • 68
  • 69
  • 70
  • 71
  • 72
  • 73
  • 74
  • 75
  • 76
  • 77
  • 78
  • 79
  • 80
  • 81
  • 82
  • 83
  • 84
  • 85
  • 86
  • 87
  • 88
  • 89
  • 90
  • 91
  • 92
  • 93
  • 94
  • 95
  • 96
  • 97
  • 98
  • 99
  • 100
  • 101
  • 102
  • 103
  • 104
  • 105
  • 106
  • 107
  • 108
  • 109
  • 110
  • 111
  • 112
  • 113
  • 114
  • 115

该程序文件是用于训练一个YOLOv5模型的。程序接受一些命令行参数,包括数据集配置文件、模型权重文件、输入图像尺寸等。训练过程中会加载模型、优化器和学习率调度器,并进行模型训练和保存权重等操作。程序还支持一些额外的功能,如模型参数冻结、模型权重的加载和保存、模型评估等。

6.系统整体结构

整体功能和构架概述:

该项目是一个基于改进的MobileNetV3骨干网络和YOLOv5的室内场景分割系统。它包含了多个程序文件,用于模型的训练、验证、推理和导出等功能。其中,train.py用于训练模型,val.py用于验证模型,ui.py用于创建用户界面,export.py用于导出模型,mobilenetv3.py和model.py分别实现了MobileNetV3和其他模型的定义和功能。

下表整理了每个文件的功能:

文件名功能
export.py将YOLOv5模型导出为其他格式
mobilenetv3.py实现MobileNetV3模型
model.py计算模型的FLOPS和参数数量
train.py训练YOLOv5模型
ui.py创建用户界面
utils.py包含一些通用的辅助函数和工具
val.py在验证数据集上验证YOLOv5模型
yolo.py实现YOLOv5模型
classify/predict.py使用分类模型进行预测
classify/train.py训练分类模型
classify/val.py在验证数据集上验证分类模型
models/common.py包含一些通用的模型组件
models/experimental.py包含一些实验性的模型
models/tf.pyTensorFlow模型相关的功能
models/yolo.pyYOLO模型相关的功能
models/init.py模型相关的初始化文件
segment/predict.py使用分割模型进行预测
segment/train.py训练分割模型
segment/val.py在验证数据集上验证分割模型
utils/activations.py包含一些激活函数
utils/augmentations.py包含一些数据增强方法
utils/autoanchor.py自动计算anchor的工具
utils/autobatch.py自动调整batch size的工具
utils/callbacks.py包含一些训练过程中的回调函数
utils/dataloaders.py包含一些数据加载器
utils/downloads.py包含一些下载相关的功能
utils/general.py包含一些通用的辅助函数
utils/loss.py包含一些损失函数
utils/metrics.py包含一些评估指标
utils/plots.py包含一些绘图相关的功能
utils/torch_utils.py包含一些PyTorch相关的工具函数
utils/triton.py与Triton Inference Server相关的功能
utils/init.py辅助函数和工具的初始化文件
utils/aws/resume.pyAWS相关的恢复功能
utils/aws/init.pyAWS相关的初始化文件
utils/flask_rest_api/example_request.pyFlask REST API的示例请求
utils/flask_rest_api/restapi.pyFlask REST API的实现
utils/loggers/init.py日志记录器的初始化文件
utils/loggers/clearml/clearml_utils.pyClearML日志记录器的工具函数
utils/loggers/clearml/hpo.pyClearML日志记录器的超参数优化功能
utils/loggers/clearml/init.pyClearML日志记录器的初始化文件
utils/loggers/comet/comet_utils.pyComet日志记录器的工具函数
utils/loggers/comet/hpo.pyComet日志记录器的超参数优化功能
utils/loggers/comet/init.pyComet日志记录器的初始化文件
utils/loggers/wandb/log_dataset.pyWandB日志记录器的数据集日志功能
utils/loggers/wandb/sweep.pyWandB日志记录器的超参数优化功能
utils/loggers/wandb/wandb_utils.pyWandB日志记录器的工具函数
utils/loggers/wandb/init.pyWandB日志记录器的初始化文件
utils/segment/augmentations.py分割模型的数据增强方法
utils/segment/dataloaders.py分割模型的数据加载器
utils/segment/general.py分割模型的通用辅助函数
utils/segment/loss.py分割模型的损失函数
utils/segment/metrics.py分割模型的评估指标
utils/segment/plots.py分割模型的绘图相关功能
utils/segment/init.py分割模型的初始化文件

请注意,由于文件数量较多,上述列表可能不完整,还可能存在其他文件。

7. MobileNetV3骨干网络

MobileNetV3 网络结构

下表给出的是 MobileNetV3-large 的网络配置。Input 表示输入当前层的特征矩阵的 shape,#out 代表的就是输出的通道大小。exp size 表示 bneck 中第一个升维的 1 × 1 1 \times 11×1 卷积输出的维度,SE 表示是否使用注意力机制,NL 表示当前使用的非线性激活函数,s 为步距 stride。bneck 后面跟的就是 DW 卷积的卷积核大小。注意最后有一个 NBN 表示分类器部分的卷积不会去使用 BN 层。

还需要注意的是第一个 bneck 结构,它的 exp size 和输出维度是一样的,也就是第一个 1 × 1 1 \times 11×1 卷积并没有做升维处理,所以在 pytorch 和 tensorflow 的官方实现中,第一个 bneck 结构中就没有使用 1 × 1 1 \times 11×1 卷积了,直接就是 DW 卷积了。

与 MobileNetV2 一致,只有当 stride = 1 且 input channel = output channel 的时候才有 shortcut 连接。

在这里插入图片描述

至于 MobileNetV3-small 的网络配置见下表,这里就不做过多赘述了。
在这里插入图片描述

更新 BlocK (bneck)

首先我们来看一下在 MobileNetV3 中 block 如何被更新的。乍一看没有太大的区别,最显眼的部分就是加入了 SE 模块,即注意力机制;其次是更新了激活函数。

在这里插入图片描述

这里的注意力机制想法非常简单,即针对每一个 channel 进行池化处理,就得到了 channel 个数个元素,通过两个全连接层,得到输出的这个向量。值得注意的是,第一个全连接层的节点个数等于 channel 个数的 1/4,然后第二个全连接层的节点就和 channel 保持一致。这个得到的输出就相当于对原始的特征矩阵的每个 channel 分析出来了其重要程度,越重要的赋予越大的权重,越不重要的就赋予越小的权重。我们用下图来进行理解,首先采用平均池化将每一个 channel 变为一个值,然后经过两个全连接层之后得到通道权重的输出,值得注意的是第二个全连接层使用 Hard-Sigmoid 激活函数。然后将通道的权重乘回原来的特征矩阵就得到了新的特征矩阵。

在这里插入图片描述

在 MobileNetV3 中 block 的激活函数标注的是 NL,表示的是非线性激活函数的意思。因为在不同层用的不一样,所以这里标注的是 NL。一样的,最后 1 × 1 1 \times 11×1 卷积后使用线性激活函数(或者说就是没有激活函数)。

重新设计激活函数

我们来重点讲一讲重新设计激活函数这个部分,之前在 MobileNetV2 都是使用 ReLU6 激活函数。现在比较常用的是 swish 激活函数,即 x 乘上 sigmoid 激活函数。使用 swish 激活函数确实能够提高网络的准确率,但是呢它也有一些问题。首先就是其计算和求导时间复杂,光一个 sigmoid 进行计算和求导就比较头疼了。第二个就是对量化过程非常不友好,特别是对于移动端的设备,为了加速一般都会进行量化操作。为此,作者提出了一个叫做 h-swish 的激活函数。
在这里插入图片描述

在说 h-swish 之前,首先要说说 h-sigmoid 激活函数,它其实是 ReLU6 ( x + 3 ) / 6 \text{ReLU6}(x+3)/6ReLU6(x+3)/6。可以看出来它和 sigmoid 非常接近,但是计算公式和求导简单太多了。由于 swish 是 x 乘上 sigmoid,自然而言得到 h-swish 是 x 乘上 h-sigmoid。可以看到 swish 激活函数的曲线和 h-swish 激活函数的曲线还是非常相似的。作者在原论文中提到,经过将 swish 激活函数替换为 h-swish,sigmoid 激活函数替换为 h-sigmoid 激活函数,对网络的推理速度是有帮助的,并且对量化过程也是很友好的。注意,h-swish 实现虽然说比 swish 快,但和 ReLU 比还是慢不少。

重新设计耗时层结构

关于重新设计耗时层结构,原论文主要讲了两个部分。首先是针对第一层卷积层,因为卷积核比较大,所以将第一层卷积核个数从 32 减少到 16。作者通过实验发现,这样做其实准确率并没有改变,但是参数量小了呀,有节省大概 2ms 的时间!
在这里插入图片描述

第二个则是精简 Last Stage。作者在使用过程中发现原始的最后结构比较耗时。精简之后第一个卷积没有变化,紧接着直接进行平均池化操作,再跟两个卷积层。和原来比起来明显少了很多层结构。作者通过实验发现这样做正确率基本没有损失,但是速度快了很多,节省了 7ms 的推理时间,别看 7ms 少,它占据了全部推理时间的 11%。

8.改进 MobileNetV3骨干网络的YOLOv5

新的层次结构或模块的引入

探索 MobileNetV3 骨干网络可能的改进空间,考虑引入新的卷积层、注意力机制或残差模块等。
例如,可以考虑添加深度可分离卷积、轻量级注意力机制或密集连接等结构,以提高网络特征表达能力或计算效率。
损失函数的优化

设计并优化适合室内场景分割任务的损失函数。

考虑综合多个指标,如像素级别准确性、边界保持能力、实时性等,设计多目标损失函数以提高模型对室内场景的分割准确度。
数据增强方法的改进

研究针对室内场景分割的数据增强策略。

考虑引入更加逼真的仿真数据、光照、旋转、尺度变化等变换,以提高模型对多样化场景的泛化能力。
对 MobileNetV3 某些层的微调或修改

分析 MobileNetV3 的结构与任务需求的匹配度,针对特定场景做出微调或修改。
可以考虑调整网络的通道数、层数或激活函数,以更好地适应室内场景分割任务的特性。

9.训练结果可视化分析

评价指标

epoch: 训练的迭代次数。
train/box_loss, train/seg_loss, train/obj_loss, train/cls_loss: 训练过程中的损失值,分别为盒子损失、分割损失、目标损失和类别损失。
metrics/precision(B), metrics/recall(B), metrics/mAP_0.5(B), metrics/mAP_0.5:0.95(B): 针对某类别(比如’B’)的精确度、召回率、mAP。
metrics/precision(M), metrics/recall(M), metrics/mAP_0.5(M), metrics/mAP_0.5:0.95(M): 另一类别(比如’M’)的对应指标。
val/box_loss, val/seg_loss, val/obj_loss, val/cls_loss: 验证集上的损失值。
x/lr0, x/lr1, x/lr2: 学习率的变化。

结果可视化

为了深入分析,我们可以从以下几个角度创建图表:

(1)训练和验证损失随着epoch的变化。
(2)不同类别的精确度、召回率和mAP随着epoch的变化。
(3)学习率的变化。

import matplotlib.pyplot as plt

# Setting up the plotting parameters
plt.figure(figsize=(15, 12))

# Plotting training losses
plt.subplot(2, 1, 1)
plt.plot(data_new['epoch'], data_new['train/box_loss'], label='Train Box Loss')
plt.plot(data_new['epoch'], data_new['train/seg_loss'], label='Train Segmentation Loss')
plt.plot(data_new['epoch'], data_new['train/obj_loss'], label='Train Object Loss')
plt.plot(data_new['epoch'], data_new['train/cls_loss'], label='Train Class Loss')
plt.title('Training Losses over Epochs')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.grid(True)

# Plotting validation losses
plt.subplot(2, 1, 2)
plt.plot(data_new['epoch'], data_new['val/box_loss'], label='Validation Box Loss')
plt.plot(data_new['epoch'], data_new['val/seg_loss'], label='Validation Segmentation Loss')
plt.plot(data_new['epoch'], data_new['val/obj_loss'], label='Validation Object Loss')
plt.plot(data_new['epoch'], data_new['val/cls_loss'], label='Validation Class Loss')
plt.title('Validation Losses over Epochs')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()

  • 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

在这里插入图片描述
现在我们有了训练和验证损失随着epoch变化的图表。接下来,我将绘制不同类别的精确度、召回率和平均精度(mAP)随着epoch的变化图表。这些指标对于理解模型在不同类别上的性能至关重要。完成这些图表后,我将进行详细的数据分析。

# Removing leading and trailing whitespaces from column names in the new dataset
data_new.columns = data_new.columns.str.strip()

# Re-plotting after correcting column names
plt.figure(figsize=(15, 12))

# Plotting training losses
plt.subplot(2, 1, 1)
plt.plot(data_new['epoch'], data_new['train/box_loss'], label='Train Box Loss')
plt.plot(data_new['epoch'], data_new['train/seg_loss'], label='Train Segmentation Loss')
plt.plot(data_new['epoch'], data_new['train/obj_loss'], label='Train Object Loss')
plt.plot(data_new['epoch'], data_new['train/cls_loss'], label='Train Class Loss')
plt.title('Training Losses over Epochs')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.grid(True)

# Plotting validation losses
plt.subplot(2, 1, 2)
plt.plot(data_new['epoch'], data_new['val/box_loss'], label='Validation Box Loss')
plt.plot(data_new['epoch'], data_new['val/seg_loss'], label='Validation Segmentation Loss')
plt.plot(data_new['epoch'], data_new['val/obj_loss'], label='Validation Object Loss')
plt.plot(data_new['epoch'], data_new['val/cls_loss'], label='Validation Class Loss')
plt.title('Validation Losses over Epochs')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.grid(True)

plt.tight_layout()
plt.show()


  • 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

在这里插入图片描述

结果数据分析
训练和验证损失(Training and Validation Losses)

损失下降趋势:训练损失(包括盒子损失、分割损失、目标损失和类别损失)随着epoch增加呈现下降趋势。这表明模型在学习过程中逐渐提高了对数据的拟合能力。

验证损失波动:验证损失的波动可能表明模型对于训练数据的学习较好,但在未见过的验证数据上的表现可能有所波动。这可能是过拟合的迹象,或者模型对新数据的泛化能力有待提高。

类别损失的特殊性:类别损失在初期下降明显,但随后趋于稳定。这可能表明模型对不同类别的分辨能力在训练初期迅速提高,后期进展放缓。

精确度和召回率(Precision and Recall)

类别B的表现:对于类别B,精确度(precision)在初期提升较快,随后趋于稳定,而召回率(recall)在整个训练过程中较为波动。高精确度表明模型能准确识别出类别B的对象,但波动的召回率可能意味着它在某些情况下错过了类别B的对象。

类别M的表现:对于类别M,精确度和召回率的变化趋势较为相似,两者都在训练初期迅速提高,随后逐渐趋于稳定。这表明模型对类别M的识别能力在初期迅速提升,后期进展放缓。

平均精度(Mean Average Precision, mAP)

类别B的mAP:类别B的mAP(在不同IoU阈值下)在训练过程中逐渐提高,说明模型在检测类别B对象的整体性能正在提升。

类别M的mAP:类别M的mAP同样显示出提升趋势,尽管其波动程度较类别B更显著。这表明模型在检测类别M时可能对某些特定情况更为敏感。

学习率(Learning Rate)

学习率调整:学习率随着训练的进行逐渐降低。这是一种常见的优化技术,称为学习率退火(learning rate annealing),旨在训练初期通过较大的学习率快速下降,并在训练后期通过较小的学习率精细调整,以避免过度拟合。

总结

整体而言,模型在训练过程中显示出了不断提升的性能,尤其在类别损失和mAP指标上有明显改善。然而,训练和验证损失的差异、以及精确度和召回率的波动,提示我们需要进一步调整模型或训练策略,以提高模型对新数据的泛化能力。具体来说,可能需要进一步的数据增强、更复杂的模型架构或更细致的参数调优。同时,观察到的过拟合迹象提示我们可能需要引入正则化技术或使用更多的数据进行训练。

10.系统整合

下图完整源码&数据集&环境部署视频教程&自定义UI界面
在这里插入图片描述

参考博客《基于改进mobilenetv3骨干网络YOLOv5的室内场景分割系统》

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/菜鸟追梦旅行/article/detail/678194
推荐阅读
相关标签
  

闽ICP备14008679号