CenterNet2實(shí)戰(zhàn)：手把手帶你實(shí)現(xiàn)使用CenterNet2訓(xùn)練自定義數(shù)據(jù)集

1、CenterNet2 介紹

論文地址：https://arxiv.org/abs/2103.07461

GitHub地址：https://github.com/xingyizhou/CenterNet2

這個(gè)模型是在Detectron2上開(kāi)發(fā)的，版本是2.3。如果是初次使用需要了解一下Detectron2。

2、Detectron2介紹

Detectron2 前身就是鼎鼎大名的 Detectron，其實(shí)Detectron可以說(shuō)是Facebook第一代檢測(cè)工具箱。Detectron2 不僅支持 Detectron已有的目標(biāo)檢測(cè)、實(shí)例分割、姿態(tài)估計(jì)等任務(wù)，還支持語(yǔ)義分割和全景分割。

優(yōu)點(diǎn)如下：

基于PyTorch：PyTorch可以提供更直觀的命令式編程模型，開(kāi)發(fā)者可以更快的進(jìn)行迭代模型設(shè)計(jì)和實(shí)驗(yàn)。

模塊化、可擴(kuò)展：從Detectron2開(kāi)始，F(xiàn)acebook引入了模塊化設(shè)計(jì)，允許用戶(hù)將自定義模塊插入目標(biāo)檢測(cè)系統(tǒng)的幾乎任何部分。這意味著許多新的研究項(xiàng)目和核心Detectron2庫(kù)可以完全分開(kāi)。其可擴(kuò)展性也使得Detectron2更加靈活。

支持語(yǔ)義分割和全景分割。

實(shí)現(xiàn)質(zhì)量：從頭開(kāi)始重寫(xiě)推出的Detectron2，使得能夠重新審視低級(jí)設(shè)計(jì)決策并解決了原始Detectron中的幾個(gè)實(shí)現(xiàn)問(wèn)題。

速度和可擴(kuò)展性：Detectron2比原始Detectron更快，而且可以更加方便進(jìn)行GPU服務(wù)器的分布式訓(xùn)練。

Detectron2go：新增了將模型產(chǎn)品化部署的軟件實(shí)現(xiàn)，包括標(biāo)準(zhǔn)的內(nèi)部數(shù)據(jù)訓(xùn)練工作流實(shí)現(xiàn)、模型壓縮量化、模型轉(zhuǎn)化等。

總之，我們使用Detectron2很方便的實(shí)現(xiàn)模型的訓(xùn)練、測(cè)試以及模型轉(zhuǎn)換。所以現(xiàn)在很多的新模型都是在Detectron2開(kāi)發(fā)。

3、搭建CenterNet2 測(cè)試環(huán)境

我本地環(huán)境：

操作系統(tǒng)：win10、Cuda11.0。

3.1 創(chuàng)建虛擬環(huán)境

創(chuàng)建虛擬環(huán)境，并激活環(huán)境。

conda create --name centernet2 python=3.7 activate centernet2 conda install pytorch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2 cudatoolkit=11.0 -c pytorch

3.2 安裝apex

APEX是英偉達(dá)開(kāi)源的，完美支持PyTorch框架，用于改變數(shù)據(jù)格式來(lái)減小模型顯存占用的工具。其中最有價(jià)值的是amp（Automatic Mixed Precision），將模型的大部分操作都用Float16數(shù)據(jù)類(lèi)型測(cè)試，一些特別操作仍然使用Float32。并且用戶(hù)僅僅通過(guò)三行代碼即可完美將自己的訓(xùn)練代碼遷移到該模型。實(shí)驗(yàn)證明，使用Float16作為大部分操作的數(shù)據(jù)類(lèi)型，并沒(méi)有降低參數(shù)，在一些實(shí)驗(yàn)中，反而由于可以增大Batch size，帶來(lái)精度上的提升，以及訓(xùn)練速度上的提升。

3.2.1 下載apex

網(wǎng)址 https://github.com/NVIDIA/apex，下載到本地文件夾。解壓后進(jìn)入到apex的目錄安裝依賴(lài)。在執(zhí)行命令；

cd C:\Users\WH\Downloads\apex-master #進(jìn)入apex目錄 pip install -r requirements.txt

3.2.2 安裝apex

依賴(lài)安裝完后，打開(kāi)cmd，cd進(jìn)入到剛剛下載完的apex-master路徑下，運(yùn)行：

python setup.py install

然后跑了一堆東西，最后是這樣的：

安裝完成！

3.3 安裝fvcore

fvcore庫(kù)的簡(jiǎn)介

fvcore是一個(gè)輕量級(jí)的核心庫(kù)，它提供了在各種計(jì)算機(jī)視覺(jué)框架(如Detectron2)中共享的最常見(jiàn)和最基本的功能。這個(gè)庫(kù)基于Python 3.6+和PyTorch。這個(gè)庫(kù)中的所有組件都經(jīng)過(guò)了類(lèi)型注釋、測(cè)試和基準(zhǔn)測(cè)試。Facebook 的人工智能實(shí)驗(yàn)室即FAIR的計(jì)算機(jī)視覺(jué)組負(fù)責(zé)維護(hù)這個(gè)庫(kù)。

github地址：https://github.com/facebookresearch/fvcore

執(zhí)行命令

conda install -c fvcore -c iopath -c conda-forge fvcore

3.4 安裝其他的庫(kù)

安裝pycocotools

pip install pycocotools

安裝cv2

pip install opencv-python

安裝 antlr4

pip install antlr4-python3-runtime

安裝future

pip install future

安裝protobuf

pip install protobuf

安裝absl

pip install absl-py

3.5 編譯CenterNet2

進(jìn)入CenterNet2目錄，目錄根據(jù)自己的實(shí)際情況更改

cd D:\CenterNet2-master

編譯

python setup.py install

4、測(cè)試環(huán)境

新建imgs和imgout文件夾，imgs文件夾存放待測(cè)試的圖片。

圖片如下：

[外鏈圖片轉(zhuǎn)存失敗,源站可能有防盜鏈機(jī)制,建議將圖片保存下來(lái)直接上傳(img-ymnqnmPL-1632828249507)(https://gitee.com/wanghao1090220084/images/raw/master/img/image-20210928095744639.png)]

執(zhí)行命令：

python projects/CenterNet2/demo.py --config-file projects/CenterNet2/configs/CenterNet2_R50_1x.yaml --input imgs/ --output imgout --opts MODEL.WEIGHTS projects/CenterNet2/CenterNet2_R50_1x.pth

運(yùn)行結(jié)果：

能夠運(yùn)行demo說(shuō)明環(huán)境已經(jīng)沒(méi)有問(wèn)題了。

5、制作數(shù)據(jù)集

本次采用的數(shù)據(jù)集是Labelme標(biāo)注的數(shù)據(jù)集，地址：鏈接：https://pan.baidu.com/s/1nxo9-NpNWKK4PwDZqwKxGQ 提取碼：kp4e，需要將其轉(zhuǎn)為COCO格式的數(shù)據(jù)集。轉(zhuǎn)換代碼如下：

新建labelme2coco.py

import argparse import json import matplotlib.pyplot as plt import skimage.io as io import cv2 from labelme import utils import numpy as np import glob import PIL.Image REQUIRE_MASK = False labels = {'aircraft': 1, 'oiltank': 2} class labelme2coco(object): def __init__(self, labelme_json=[], save_json_path='./new.json'): ''' :param labelme_json: the list of all labelme json file paths :param save_json_path: the path to save new json ''' self.labelme_json = labelme_json self.save_json_path = save_json_path self.images = [] self.categories = [] self.annotations = [] # self.data_coco = {} self.label = [] self.annID = 1 self.height = 0 self.width = 0 self.require_mask = REQUIRE_MASK self.save_json() def data_transfer(self): for num, json_file in enumerate(self.labelme_json): if not json_file == self.save_json_path: with open(json_file, 'r') as fp: data = json.load(fp) self.images.append(self.image(data, num)) for shapes in data['shapes']: print("label is ") print(shapes['label']) label = shapes['label'] # if label[1] not in self.label: if label not in self.label: print("find new category: ") self.categories.append(self.categorie(label)) print(self.categories) # self.label.append(label[1]) self.label.append(label) points = shapes['points'] self.annotations.append(self.annotation(points, label, num)) self.annID += 1 def image(self, data, num): image = {} img = utils.img_b64_to_arr(data['imageData']) height, width = img.shape[:2] img = None image['height'] = height image['width'] = width image['id'] = num + 1 image['file_name'] = data['imagePath'].split('/')[-1] self.height = height self.width = width return image def categorie(self, label): categorie = {} categorie['supercategory'] = label # categorie['supercategory'] = label categorie['id'] = labels[label] # 0 默認(rèn)為背景 categorie['name'] = label return categorie def annotation(self, points, label, num): annotation = {} print(points) x1 = points[0][0] y1 = points[0][1] x2 = points[1][0] y2 = points[1][1] contour = np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]]) # points = [[x1, y1], [x2, y2]] for rectangle contour = contour.astype(int) area = cv2.contourArea(contour) print("contour is ", contour, " area = ", area) annotation['segmentation'] = [list(np.asarray([[x1, y1], [x2, y1], [x2, y2], [x1, y2]]).flatten())] # [list(np.asarray(contour).flatten())] annotation['iscrowd'] = 0 annotation['area'] = area annotation['image_id'] = num + 1 if self.require_mask: annotation['bbox'] = list(map(float, self.getbbox(points))) else: x1 = points[0][0] y1 = points[0][1] width = points[1][0] - x1 height = points[1][1] - y1 annotation['bbox'] = list(np.asarray([x1, y1, width, height]).flatten()) annotation['category_id'] = self.getcatid(label) annotation['id'] = self.annID return annotation def getcatid(self, label): for categorie in self.categories: # if label[1]==categorie['name']: if label == categorie['name']: return categorie['id'] return -1 def getbbox(self, points): polygons = points mask = self.polygons_to_mask([self.height, self.width], polygons) return self.mask2box(mask) def mask2box(self, mask): # np.where(mask==1) index = np.argwhere(mask == 1) rows = index[:, 0] clos = index[:, 1] left_top_r = np.min(rows) # y left_top_c = np.min(clos) # x right_bottom_r = np.max(rows) right_bottom_c = np.max(clos) return [left_top_c, left_top_r, right_bottom_c - left_top_c, right_bottom_r - left_top_r] def polygons_to_mask(self, img_shape, polygons): mask = np.zeros(img_shape, dtype=np.uint8) mask = PIL.Image.fromarray(mask) xy = list(map(tuple, polygons)) PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1) mask = np.array(mask, dtype=bool) return mask def data2coco(self): data_coco = {} data_coco['images'] = self.images data_coco['categories'] = self.categories data_coco['annotations'] = self.annotations return data_coco def save_json(self): print("in save_json") self.data_transfer() self.data_coco = self.data2coco() print(self.save_json_path) json.dump(self.data_coco, open(self.save_json_path, 'w'), indent=4) labelme_json = glob.glob('LabelmeData/*.json') from sklearn.model_selection import train_test_split trainval_files, test_files = train_test_split(labelme_json, test_size=0.2, random_state=55) import os if not os.path.exists("projects/CenterNet2/datasets/coco/annotations"): os.makedirs("projects/CenterNet2/datasets/coco/annotations/") if not os.path.exists("projects/CenterNet2/datasets/coco/train2017"): os.makedirs("projects/CenterNet2/datasets/coco/train2017") if not os.path.exists("projects/CenterNet2/datasets/coco/val2017"): os.makedirs("projects/CenterNet2/datasets/coco/val2017") labelme2coco(trainval_files, 'projects/CenterNet2/datasets/coco/annotations/instances_train2017.json') labelme2coco(test_files, 'projects/CenterNet2/datasets/coco/annotations/instances_val2017.json') import shutil for file in trainval_files: shutil.copy(os.path.splitext(file)[0] + ".jpg", "projects/CenterNet2/datasets/coco/train2017/") for file in test_files: shutil.copy(os.path.splitext(file)[0] + ".jpg", "projects/CenterNet2/datasets/coco/val2017/")

6、配置訓(xùn)練環(huán)境

6.1 更改預(yù)訓(xùn)練模型的size

在projects/CenterNet2目錄，新建change_model_size.py文件

import torch import numpy as np import pickle num_class = 2 pretrained_weights = torch.load('models/CenterNet2_R50_1x.pth') pretrained_weights['iteration']=0 pretrained_weights['model']["roi_heads.box_predictor.0.cls_score.weight"].resize_(num_class+1,1024) pretrained_weights['model']["roi_heads.box_predictor.0.cls_score.bias"].resize_(num_class+1) pretrained_weights['model']["roi_heads.box_predictor.1.cls_score.weight"].resize_(num_class+1,1024) pretrained_weights['model']["roi_heads.box_predictor.1.cls_score.bias"].resize_(num_class+1) pretrained_weights['model']["roi_heads.box_predictor.2.cls_score.weight"].resize_(num_class+1,1024) pretrained_weights['model']["roi_heads.box_predictor.2.cls_score.bias"].resize_(num_class+1) torch.save(pretrained_weights, "models/CenterNet2_%d.pth"%num_class)

這個(gè)文件的目的是修改模型輸出的size，numclass按照本次打算訓(xùn)練的數(shù)據(jù)集的類(lèi)別設(shè)置。

6.2 修改config參數(shù)

路徑：“detectron2/engine/defaults.py”

–config-file：模型的配置文件，CenterNet2的模型配置文件放在“projects/CenterNet2/configs”下面。名字和預(yù)訓(xùn)練模型對(duì)應(yīng)。

parser.add_argument("--config-file", default="./configs/CenterNet2_DLA-BiFPN-P3_4x.yaml", metavar="FILE", help="path to config file")

resume 是否再次，訓(xùn)練，如果設(shè)置為true，則接著上次訓(xùn)練的結(jié)果訓(xùn)練。所以第一次訓(xùn)練不用設(shè)置。

parser.add_argument( "--resume", action="store_true", help="Whether to attempt to resume from the checkpoint directory. " "See documentation of `DefaultTrainer.resume_or_load()` for what it means.", )

–num-gpus，gpu的個(gè)數(shù)，如果只有一個(gè)設(shè)置為1，如果有多個(gè)，可以自己設(shè)置想用的個(gè)數(shù)。

parser.add_argument("--num-gpus", type=int, default=1, help="number of gpus *per machine*")

opts指的是yaml文件的參數(shù)。

上面的參數(shù)可以設(shè)置，也可以不設(shè)置，設(shè)置之后可以直接運(yùn)行不用再考慮設(shè)置參數(shù)，如果不設(shè)置每次訓(xùn)練的時(shí)候配置一次參數(shù)。

修改類(lèi)別，文件路徑“projects/CenterNet2/centernet/config.py”，

_C.MODEL.CENTERNET.NUM_CLASSES = 2

修改yaml文件參數(shù)

Base-CenterNet2.yaml中修改預(yù)訓(xùn)練模型的路徑。

WEIGHTS: "CenterNet2_2.pth"

BASE_LR：設(shè)置學(xué)習(xí)率。

STEPS：設(shè)置訓(xùn)練多少步之后調(diào)整學(xué)習(xí)率。

MAX_ITER：最大迭代次數(shù)。

CHECKPOINT_PERIOD：設(shè)置迭代多少次保存一次模型

BASE_LR: 0.01 STEPS: (10000, 50000) MAX_ITER: 100000 CHECKPOINT_PERIOD: 5000

在設(shè)置上面的參數(shù)時(shí)要注意，如果選擇用CenterNet2_R50_1x.yaml，里面沒(méi)有參數(shù)，則在Base-CenterNet2.yaml中設(shè)置，如果選用其他的，例如CenterNet2_DLA-BiFPN-P3_4x.yaml，這些參數(shù)需要在CenterNet2_DLA-BiFPN-P3_4x.yaml改。

6.3 修改train_net.py

主要修改該setup函數(shù)，增加數(shù)據(jù)集注冊(cè)。

NUM_CLASSES=2 def setup(args): """ Create configs and perform basic setups. """ register_coco_instances("train", {}, "datasets/coco/annotations/instances_train2017.json", "datasets/coco/train2017") register_coco_instances("test", {}, "datasets/coco/annotations/instances_val2017.json", "datasets/coco/val2017") cfg = get_cfg() add_centernet_config(cfg) cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.DATASETS.TRAIN = ("train",) cfg.DATASETS.TEST = ("test",) cfg.MODEL.CENTERNET.NUM_CLASSES = NUM_CLASSES cfg.MODEL.ROI_HEADS.NUM_CLASSES = NUM_CLASSES if '/auto' in cfg.OUTPUT_DIR: file_name = os.path.basename(args.config_file)[:-5] cfg.OUTPUT_DIR = cfg.OUTPUT_DIR.replace('/auto', '/{}'.format(file_name)) logger.info('OUTPUT_DIR: {}'.format(cfg.OUTPUT_DIR)) cfg.freeze() default_setup(cfg, args) return cfg

還要修改detectron2/engine/launch.py，在launch函數(shù)下面增加一句

dist.init_process_group('gloo', init_method='file://tmp/somefile', rank=0, world_size=1)

如下圖：

這句話(huà)的作用是初始化分布式訓(xùn)練，因?yàn)槲覀儧](méi)有使用分布式，所以沒(méi)有初始化，但是不初始化就會(huì)報(bào)錯(cuò)，所以加上這句。

7、訓(xùn)練

兩種啟動(dòng)方式：

第一種，命令行：進(jìn)入“projects/CenterNet2/”目錄下,執(zhí)行：

python train_net.py

第二種，直接在pycharm 直接運(yùn)行train_net.py.

訓(xùn)練結(jié)果：

從訓(xùn)練結(jié)果上看，效果確實(shí)不錯(cuò)，不過(guò)模型很大。大約有500M

8、測(cè)試

修改projects/CenterNet2/demo.py

8.1 修改setup_cfg函數(shù)

在紅框的位置增加代碼，詳細(xì)如下面的代碼。

NUM_CLASSES=2 def setup_cfg(args): # load config from file and command-line arguments cfg = get_cfg() add_centernet_config(cfg) cfg.MODEL.CENTERNET.NUM_CLASSES = NUM_CLASSES cfg.MODEL.ROI_HEADS.NUM_CLASSES = NUM_CLASSES cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) # Set score_threshold for builtin models cfg.MODEL.RETINANET.SCORE_THRESH_TEST = args.confidence_threshold cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = args.confidence_threshold if cfg.MODEL.META_ARCHITECTURE in ['ProposalNetwork', 'CenterNetDetector']: cfg.MODEL.CENTERNET.INFERENCE_TH = args.confidence_threshold cfg.MODEL.CENTERNET.NMS_TH = cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST cfg.MODEL.PANOPTIC_FPN.COMBINE.INSTANCES_CONFIDENCE_THRESH = args.confidence_threshold cfg.freeze() return cfg

8.2 修改顯示類(lèi)別

代碼：

visualizer.metadata.thing_classes[:10] = ["aircraft", "oiltank"]

然后進(jìn)入CenterNet2-master目錄，執(zhí)行如下命令：

python projects/CenterNet2/demo.py --config-file projects/CenterNet2/configs/CenterNet2_R50_1x.yaml --input imgs/ --output imgout --opts MODEL.WEIGHTS projects/CenterNet2/output/CenterNet2/CenterNet2_R50_1x/model_final.pth

運(yùn)行結(jié)果：

關(guān)注公眾號(hào)“AI小浩”，回復(fù)“centernet2實(shí)戰(zhàn)”，獲取本次項(xiàng)目源碼和本文的PDF版本。

機(jī)器學(xué)習(xí) 深度學(xué)習(xí)

版權(quán)聲明：本文內(nèi)容由網(wǎng)絡(luò)用戶(hù)投稿，版權(quán)歸原作者所有，本站不擁有其著作權(quán)，亦不承擔(dān)相應(yīng)法律責(zé)任。如果您發(fā)現(xiàn)本站中有涉嫌抄襲或描述失實(shí)的內(nèi)容，請(qǐng)聯(lián)系我們jiasou666@gmail.com 處理，核實(shí)后本網(wǎng)站將在24小時(shí)內(nèi)刪除侵權(quán)內(nèi)容。