1.安装mmengine和mmcv依赖
pip install -U openmim
pip install "mmengine=0.7.1"
mim install "mmcv==2.0.0"
2.安装mmdetection模块
git clone -b tutorials https://github.com/open-mmlab/mmdetection.git
3.准备数据集
#数据集下载
!rm -rf cat_dataset
!wget https://download.openmmlab.com/mmyolo/data/balloon_dataset.zip
!unzip balloon_dataset -d balloon_dataset && rm balloon_dataset.zip
4.数据集可视化
import matplotlib.pyplot as plt
from PIL import Image
%matplotlib inline
%config InlineBackend.figure_format = 'retina'
original_images = []
images = []
texts = []
plt.figure(figsize=(16, 5))
image_paths = [filename for filename in os.listdir('balloon_dataset/balloon/train')][:8]
for i, filename in enumerate(image_paths):
name = os.path.splitext(filename)[0]
image = Image.open('balloon_dataset/balloon/train/' + filename).convert("RGB")
plt.subplot(2, 4, i+1)
plt.imshow(image)
plt.title(f"{filename}")
plt.xticks([])
plt.yticks([])
plt.tight_layout()
5.修改配置参数
config_balloon = """
# Inherit and overwrite part of the config based on this config
_base_ = './rtmdet_tiny_8xb32-300e_coco.py'
data_root = 'data/balloon/' # dataset root
train_batch_size_per_gpu = 4
train_num_workers = 2
max_epochs = 20
stage2_num_epochs = 1
base_lr = 0.00008
metainfo = {
'classes': ('balloon', ),
'palette': [
(220, 20, 60),
]
}
train_dataloader = dict(
batch_size=train_batch_size_per_gpu,
num_workers=train_num_workers,
dataset=dict(
data_root=data_root,
metainfo=metainfo,
data_prefix=dict(img='train/'),
ann_file='train.json'))
val_dataloader = dict(
dataset=dict(
data_root=data_root,
metainfo=metainfo,
data_prefix=dict(img='val/'),
ann_file='val.json'))
test_dataloader = val_dataloader
val_evaluator = dict(ann_file=data_root + 'val.json')
test_evaluator = val_evaluator
model = dict(bbox_head=dict(num_classes=1))
# learning rate
param_scheduler = [
dict(
type='LinearLR',
start_factor=1.0e-5,
by_epoch=False,
begin=0,
end=10),
dict(
# use cosine lr from 10 to 20 epoch
type='CosineAnnealingLR',
eta_min=base_lr * 0.05,
begin=max_epochs // 2,
end=max_epochs,
T_max=max_epochs // 2,
by_epoch=True,
convert_to_iter_based=True),
]
train_pipeline_stage2 = [
dict(type='LoadImageFromFile', backend_args=None),
dict(type='LoadAnnotations', with_bbox=True),
dict(
type='RandomResize',
scale=(640, 640),
ratio_range=(0.1, 2.0),
keep_ratio=True),
dict(type='RandomCrop', crop_size=(640, 640)),
dict(type='YOLOXHSVRandomAug'),
dict(type='RandomFlip', prob=0.5),
dict(type='Pad', size=(640, 640), pad_val=dict(img=(114, 114, 114))),
dict(type='PackDetInputs')
]
# optimizer
optim_wrapper = dict(
_delete_=True,
type='OptimWrapper',
optimizer=dict(type='AdamW', lr=base_lr, weight_decay=0.05),
paramwise_cfg=dict(
norm_decay_mult=0, bias_decay_mult=0, bypass_duplicate=True))
default_hooks = dict(
checkpoint=dict(
interval=5,
max_keep_ckpts=2, # only keep latest 2 checkpoints
save_best='auto'
),
logger=dict(type='LoggerHook', interval=5))
custom_hooks = [
dict(
type='PipelineSwitchHook',
switch_epoch=max_epochs - stage2_num_epochs,
switch_pipeline=train_pipeline_stage2)
]
# load COCO pre-trained weight
load_from = './checkpoints/rtmdet_tiny_8xb32-300e_coco_20220902_112414-78e30dcc.pth'
train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=max_epochs, val_interval=1)
visualizer = dict(vis_backends=[dict(type='LocalVisBackend'),dict(type='TensorboardVisBackend')])
"""
with open('./configs/rtmdet/rtmdet_tiny_1xb4-20e_balloon.py', 'w') as f:
f.write(config_balloon)
6.训练
python tools/train.py configs/rtmdet/rtmdet_tiny_1xb4-20e_balloon.py
7.测试
python tools/test.py configs/rtmdet/rtmdet_tiny_1xb4-20e_balloon.py work_dirs/rtmdet_tiny_1xb4-20e_balloon/best_coco_bbox_mAP_epoch_19.pth