🏖️
mmdetectionをmmdeployでonnx
TL;DR
mmdetとmmdeployで以下の手順で学習と変換
$ python tools/train.py \
mmdet_config_model.py # configから学習モデルを作成
$ python mmdeploy/tools/deploy.py \
mmdeploy_config_detection_ascend_static-640x640.py.py \ # mmdeployのconfig
mmdet_config_model.py \ # mmdetモデルのconfig
checkpoint.pth \ # 学習済みモデル
demo_input.png \ # ダミー入力
--work-dir mmdeploy_model/rtmdet # onnxモデルの保存先
奇っ怪なタイトルですが、行うことはシンプルです。
- データセットDocLayNetをフレームワークmmdetectionでモデルRTMDetを学習させる
- 学習したRTMDetをフレームワークmmdeployでonnxファイルに変換する
この記事では、初めて使う人でもとりあえず動かせるように、コード全体をペタペタ貼り付けて書いて行きます。
DocLayNet はPDFドキュメントを画像にしたものと、文字・表・タイトルなどカテゴリでアノテーションされたCOCOフォーマットのデータセットです。物体検出をつかったレイアウト検出や、BERTなどの系列処理を用いたレイアウト提案などで用いられます。
今回は、Apache License 2.0でライセンシングされたYOLO系の物体検出モデルRTMDet を使ってDocLayNetを学習し、ドキュメント画像からコンテンツをBBoxで囲むようなものを作ります。
さらにこれをONNX (Open Neural Network Exchange)フォーマットに変換することで、外部で使いやすい形で保存することができます。
ここで使用するのがmmdetection とmmdeploy です。
mmdetectionは物体検出モデルやinstance segmentationモデルを収録し、さらに学習関連の最適化、オーグメンテーション、可視化など非常に多くのことが可能なライブラリです。
RTMDetもこれに収録されていて、pythonの辞書を用いて書かれるconfigファイルを弄ることでモデルとデータローダー、最適化の一連の流れを組み立てられます。
mmdeployはmm open lab系のライブラリで作ったモデルを様々な形式に変換することができるライブラリです。今回はonnxへ変換しますが、他にもNvidia系のGPUで高速に動作するTensorRTや、iPhone上で動作するCoreMLなどにも変換できます。
mmdetectionの環境構築
公式のDockerfileをbuildするだけですぐ使えます。
このコンテナに、予め用意したデータセットなどが置いてあるディレクトリをマウントして、以下のようなディレクトリ構成を作ります。構成自体は自分が使いやすいように適当にしても動きます。
/workspaces/doclaynet/
├── dataset/
│ ├── COCO/ ... train.json test.json
│ └── PNG/ ... たくさんのimg.png
└── mmdetection/(symbolic link)
├── _myconfig/ ... 自作のrtmdet_s.py
├── _other/ ... git cloneでついてきたいらないディレクトリやファイル
├── work_dir/ ... 学習したlogやchekpoint(モデル本体)が格納される
├── configs/ ... mmdetectionのconfig一覧
└── tools/ ... train.py test.py など
mmdetectionはもとあるディレクトリからリンクを貼っています。ln -s {つなげたいdir} ./{リンクの名前} とすれば現在いる所にシンボリックリンクを作れます。
使わないものは_other/にまとめました。
configファイルは_myconfig/に置いていきます。公式では元のディレクトリに継承して書いて行くことを推奨していますが、私はこのほうが実験管理が楽なので、継承すら使わずに全分手動コピーして書いています。
config作成と学習
以下のconfigはrtmdet-sの以下のファイルをDoclaynet用に改造し、全てのパラメータをスクラッチで書いたものです。
改造した部分はNOTEでコメントしてあります。
実際は値を直書きせず、上にコメントアウトしてあるような自分用変数を用意して、それを使う運用が良いです。
rtmdet_s_doclaynet.py
# 公式のCOCO Object Detection学習済みの重み
# BACKBONE_WEIGHT = 'https://download.openmmlab.com/mmdetection/v3.0/rtmdet/cspnext_rsb_pretrain/cspnext-s_imagenet_600e.pth'
# データセットPNG/とCOCO/が置いてあるディレクトリ
# BASEDIR_DATASET = '/workspaces/doclaynet/dataset/'
# DATA_IMAGE = 'PNG/'
# ANNOFILE_TRAIN = 'COCO/train.json'
# ANNOFILE_TEST = 'COCO/test.json'
# DocLayNetのクラス(COCOのフォーマットだが、カテゴリなどが異なる場合はこちらで設定してあげる必要がある 忘れがちなので注意)
# CLASSES = ('CaptionFootnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title')
# NUM_CLASS = 10
# LOG_INTERVAL = 100 # 通知するタイミング(処理する枚数)
# BATCHSIZE = 4 # mini batch 枚数
# LEARNING_RATE = 0.001 # 最大学習率(オリジナルのRTMDetでは0.004)
# WEIGHT_DCAY = 0.01 # weight decay(オリジナルではもうちょっとでかい)
# MAX_EPOCH = 16 # 学習epoch数(オリジナルでは300)
# 物体検出モデル本体の設定
model = dict(
type='RTMDet',
data_preprocessor=dict(
type='DetDataPreprocessor',
mean=[103.53, 116.28, 123.675],
std=[57.375, 57.12, 58.395],
bgr_to_rgb=False,
batch_augments=None),
backbone=dict(
type='CSPNeXt',
arch='P5',
expand_ratio=0.5,
deepen_factor=0.33,
widen_factor=0.5,
channel_attention=True,
norm_cfg=dict(type='SyncBN'),
act_cfg=dict(type='SiLU', inplace=True),
init_cfg=dict(
type='Pretrained',
prefix='backbone.',
checkpoint=
'https://download.openmmlab.com/mmdetection/v3.0/rtmdet/cspnext_rsb_pretrain/cspnext-s_imagenet_600e.pth'
)),
neck=dict(
type='CSPNeXtPAFPN',
in_channels=[128, 256, 512],
out_channels=128,
num_csp_blocks=1,
expand_ratio=0.5,
norm_cfg=dict(type='SyncBN'),
act_cfg=dict(type='SiLU', inplace=True)),
bbox_head=dict(
type='RTMDetSepBNHead',
num_classes=10, # NOTE
in_channels=128,
feat_channels=128,
stacked_convs=2,
anchor_generator=dict(
type='MlvlPointGenerator', offset=0, strides=[8, 16, 32]),
bbox_coder=dict(type='DistancePointBBoxCoder'),
loss_cls=dict(
type='QualityFocalLoss',
use_sigmoid=True,
beta=2.0,
loss_weight=1.0),
loss_bbox=dict(type='GIoULoss', loss_weight=2.0),
with_objectness=False,
exp_on_reg=False,
share_conv=True,
pred_kernel_size=1,
norm_cfg=dict(type='SyncBN'),
act_cfg=dict(type='SiLU', inplace=True)),
train_cfg=dict(
assigner=dict(type='DynamicSoftLabelAssigner', topk=2), # NOTE ソフトラベルを作成するtop-kはクラス数内の値
allowed_border=-1,
pos_weight=-1,
debug=False),
test_cfg=dict(
nms_pre=30000,
min_bbox_size=0,
score_thr=0.001,
nms=dict(type='nms', iou_threshold=0.65),
max_per_img=300))
# データセットとデータローダー、データオーグメンテーションの設定
train_pipeline = [ # NOTE 今回PDFドキュメントなので、変な回転や切り抜きは行わない
dict(type='LoadImageFromFile', backend_args=None),
dict(type='LoadAnnotations', with_bbox=True),
dict(type='Resize', scale=(640, 640), keep_ratio=True),
dict(type='YOLOXHSVRandomAug'),
dict(type='RandomFlip', prob=0.2),
dict(type='Pad', size=(640, 640), pad_val=dict(img=(114, 114, 114))),
dict(type='PackDetInputs')
]
test_pipeline = [
dict(type='LoadImageFromFile', backend_args=None),
dict(type='LoadAnnotations', with_bbox=True),
dict(type='Resize', scale=(640, 640), keep_ratio=True),
dict(type='Pad', size=(640, 640), pad_val=dict(img=(114, 114, 114))),
dict(
type='PackDetInputs',
meta_keys=('img_id', 'img_path', 'ori_shape', 'img_shape',
'scale_factor'))
]
train_dataset = dict(
type='CocoDataset',
data_root='/workspaces/doclaynet/dataset/', # NOTE
ann_file='COCO/train.json', # NOTE
data_prefix=dict(img='PNG/'), # NOTE
filter_cfg=dict(filter_empty_gt=True, min_size=32),
pipeline=train_pipline,
metainfo=dict(
classes=('CaptionFootnote', 'Formula', 'List-item', 'Page-footer',
'Page-header', 'Picture', 'Section-header', 'Table', 'Text',
'Title')),
backend_args=None)
test_dataset = dict(
type='CocoDataset',
data_root='/workspaces/doclaynet/dataset/',
ann_file='COCO/test.json',
data_prefix=dict(img='PNG/'),
test_mode=True,
pipeline=test_pipline,
metainfo=dict(
classes=('CaptionFootnote', 'Formula', 'List-item', 'Page-footer',
'Page-header', 'Picture', 'Section-header', 'Table', 'Text',
'Title')),
backend_args=None)
train_dataloader = dict(
batch_size=4, # NOTE
num_workers=2,
persistent_workers=True,
sampler=dict(type='DefaultSampler', shuffle=True),
batch_sampler=None,
pin_memory=True,
dataset=train_dataset)
val_dataloader = dict(
batch_size=1,
num_workers=2,
persistent_workers=True,
drop_last=False,
sampler=dict(type='DefaultSampler', shuffle=False),
dataset=test_dataset)
test_dataloader = dict(
batch_size=1,
num_workers=2,
persistent_workers=True,
drop_last=False,
sampler=dict(type='DefaultSampler', shuffle=False),
dataset=test_dataset)
# 学習と検証のロジック
val_evaluator = dict(
type='CocoMetric',
ann_file='/workspaces/doclaynet/dataset/COCO/test.json', # NOTE
metric='bbox',
format_only=False,
backend_args=None,
proposal_nums=(100, 1, 10))
test_evaluator = dict(
type='CocoMetric',
ann_file='/workspaces/doclaynet/dataset/COCO/test.json', # NOTE
metric='bbox',
format_only=False,
backend_args=None,
proposal_nums=(100, 1, 10))
train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=16, val_interval=100) # NOTE
val_cfg = dict(type='ValLoop')
test_cfg = dict(type='TestLoop')
base_lr = 0.001 # NOTE
# NOTE 最適化は単純なcosine annealingだけにした(オリジナルはwarmup付き)
optim_wrapper = dict(
type='OptimWrapper',
optimizer=dict(type='AdamW', lr=0.001, weight_decay=0.01), # NOTE
paramwise_cfg=dict(
norm_decay_mult=0, bias_decay_mult=0, bypass_duplicate=True))
param_scheduler = dict(
type='CosineAnnealingLR',
eta_min=1e-05,
begin=0,
end=16,
T_max=16,
by_epoch=True,
convert_to_iter_based=True)
# その他logingなどの設定
default_scope = 'mmdet'
default_hooks = dict(
timer=dict(type='IterTimerHook'),
logger=dict(type='LoggerHook', interval=100), # NOTE
param_scheduler=dict(type='ParamSchedulerHook'),
checkpoint=dict(type='CheckpointHook', interval=1),
sampler_seed=dict(type='DistSamplerSeedHook'),
visualization=dict(type='DetVisualizationHook'))
custom_hooks = ({
'type': 'EMAHook',
'ema_type': 'ExpMomentumEMA',
'momentum': 0.0002,
'update_buffers': True,
'priority': 49
}, )
env_cfg = dict(
cudnn_benchmark=False,
mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
dist_cfg=dict(backend='nccl'))
vis_backends = [dict(type='LocalVisBackend')]
visualizer = dict(
type='DetLocalVisualizer',
vis_backends=[dict(type='LocalVisBackend')],
name='visualizer')
log_processor = dict(type='LogProcessor', window_size=100, by_epoch=True)
log_level = 'INFO'
load_from = None
resume = False
launcher = 'none'
work_dir = './work_dirs/rtmdet_s'
これでRTMDetの設定は終了で、tools/train.pyを実行してモデルを学習する。
python tools/train.py /workspaces/doclaynet/mmdetection/_myconfig/rtmdet_s.py
これにより以下のようなlogが長々と表示され、学習が終わるまで1日ちょっと待つ。
System environment:
sys.platform: linux
Python: 3.7.10 (default, Feb 26 2021, 18:47:35) [GCC 7.3.0]
CUDA available: True
numpy_random_seed: 1683244977
GPU 0: Quadro P5000
CUDA_HOME: /usr/local/cuda
NVCC: Cuda compilation tools, release 11.1, V11.1.105
GCC: gcc (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
PyTorch: 1.9.0
PyTorch compiling details: PyTorch built with:
- GCC 7.3
- C++ Version: 201402
- Intel(R) oneAPI Math Kernel Library Version 2021.2-Product Build 20210312 for Intel(R) 64 architecture applications
- Intel(R) MKL-DNN v2.1.2 (Git Hash 98be7e8afa711dc9b66c8ff3504129cb82013cdb)
- OpenMP 201511 (a.k.a. OpenMP 4.5)
- NNPACK is enabled
- CPU capability usage: AVX2
- CUDA Runtime 11.1
- NVCC architecture flags: -gencode;arch=compute_37,code=sm_37;-gencode;arch=compute_50,code=sm_50;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_61,code=sm_61;-gencode;arch=compute_70,code=sm_70;-gencode;arch=compute_75,code=sm_75;-gencode;arch=compute_80,code=sm_80;-gencode;arch=compute_86,code=sm_86;-gencode;arch=compute_37,code=compute_37
- CuDNN 8.0.5
- Magma 2.5.2
- Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CUDA_VERSION=11.1, CUDNN_VERSION=8.0.5, CXX_COMPILER=/opt/rh/devtoolset-7/root/usr/bin/c++, CXX_FLAGS= -Wno-deprecated -fvisibility-inlines-hidden -DUSE_PTHREADPOOL -fopenmp -DNDEBUG -DUSE_KINETO -DUSE_FBGEMM -DUSE_QNNPACK -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -Wno-narrowing -Wall -Wextra -Werror=return-type -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-sign-compare -Wno-unused-parameter -Wno-unused-variable -Wno-unused-function -Wno-unused-result -Wno-unused-local-typedefs -Wno-strict-overflow -Wno-strict-aliasing -Wno-error=deprecated-declarations -Wno-stringop-overflow -Wno-psabi -Wno-error=pedantic -Wno-error=redundant-decls -Wno-error=old-style-cast -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Werror=format -Wno-stringop-overflow, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_VERSION=1.9.0, USE_CUDA=ON, USE_CUDNN=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, USE_NCCL=ON, USE_NNPACK=ON, USE_OPENMP=ON,
TorchVision: 0.10.0
OpenCV: 4.7.0
MMEngine: 0.7.3
Runtime environment:
cudnn_benchmark: False
mp_cfg: {'mp_start_method': 'fork', 'opencv_num_threads': 0}
dist_cfg: {'backend': 'nccl'}
seed: None
Distributed launcher: none
Distributed training: False
GPU number: 1
~
Initialized by user-defined `init_weights` in RTMDetSepBNHead
2023/05/15 12:29:44 - mmengine - INFO - Load checkpoint from /workspaces/doclaynet/mmdetection/work_dirs/rtmdet_s/epoch_12.pth
2023/05/15 12:29:44 - mmengine - INFO - resumed epoch: 12, iter: 207312
2023/05/15 12:29:44 - mmengine - WARNING - "FileClient" will be deprecated in future. Please use io functions in https://mmengine.readthedocs.io/en/latest/api/fileio.html#file-io
2023/05/15 12:29:44 - mmengine - WARNING - "HardDiskBackend" is the alias of "LocalBackend" and the former will be deprecated in future.
2023/05/15 12:29:44 - mmengine - INFO - Checkpoints will be saved to /mmdetection/work_dirs/rtmdet_s.
2023/05/14 18:10:15 - mmengine - INFO - Epoch(train) [1][ 100/17276] lr: 1.0000e-03 eta: 1 day, 3:22:28 time: 0.3567 data_time: 0.0061 memory: 2006 loss: 0.9077 loss_cls: 0.5126 loss_bbox: 0.3951
2023/05/14 18:10:48 - mmengine - INFO - Epoch(train) [1][ 200/17276] lr: 1.0000e-03 eta: 1 day, 2:13:41 time: 0.3270 data_time: 0.0011 memory: 2004 loss: 1.0106 loss_cls: 0.5878 loss_bbox: 0.4228
2023/05/14 18:11:14 - mmengine - INFO - Epoch(train) [1][ 300/17276] lr: 1.0000e-03 eta: 1 day, 0:13:27 time: 0.2638 data_time: 0.0011 memory: 1983 loss: 0.9623 loss_cls: 0.5487 loss_bbox: 0.4136
~
2023/05/15 18:40:53 - mmengine - INFO - Epoch(train) [16][17200/17276] lr: 1.0000e-05 eta: 0:00:24 time: 0.2662 data_time: 0.0017 memory: 1988 loss: 0.3761 loss_cls: 0.1775 loss_bbox: 0.1986
2023/05/15 18:41:15 - mmengine - INFO - Exp name: rtmdet_s_20230515_122852
2023/05/15 18:41:15 - mmengine - INFO - Saving checkpoint at 16 epochs
データの可視化
作られたモデルの推論結果などは次のスクリプトで見ることができます。
# データ可視化
(doclaynet/)
python mmdetection/tools/test.py {config.py} {checkpoint.pth} --show-dir vis/
input image
バグで正解データがちょっとずれているが、大体あっているのでOKです。
左がずれた正解データ、右が推論結果
# loss可視化
python tools/analysis_tools/analyze_logs.py plot_curve {log.json} --keys loss_cls loss_bbox loss --out losses.png
途中事故で学習が止まったのでグロッキングっぽい挙動になっているが、普通は滑らかに学習されるはず...
train loss
mmdeployの環境設定
こちらも配布されているDockerを使います。
おそらくmmdetectionのコンテナを使うよりこっちを独立させたほうが楽(何も考えなくてよいので)
onnxへ変換
次のコードはdetection_ascend_static-640x640.pyの展開。
今回使ったRTMDetの学習時のinputの大きさが640x640pxなので、それを静的推論版を選びます。
onnx_config = dict(
type='onnx',
export_params=True,
keep_initializers_as_inputs=False,
opset_version=11,
save_file='end2end.onnx',
input_names=['input'],
output_names=['dets', 'labels'],
input_shape=[640, 640],
optimize=True)
codebase_config = dict(
type='mmdet',
task='ObjectDetection',
model_type='end2end',
post_processing=dict(
score_threshold=0.05,
confidence_threshold=0.005,
iou_threshold=0.5,
max_output_boxes_per_class=200,
pre_top_k=5000,
keep_top_k=100,
background_label_id=-1,
))
backend_config = dict(
type='ascend',
model_inputs=[dict(input_shapes=dict(input=[1, 3, 640, 640]))]
)
これを次のスクリプトで実行。
この前に mim install mmdet でpythonからmmdetを使えるようにするのを忘れずに。
python mmdeploy/tools/deploy.py \
# mmdeployのconfig
/workspaces/doclaynet_deploy/mmdeploy/_myconfig/det_to_onnx.py \
# mmdetectionのconfig
/workspaces/doclaynet_deploy/rtmdet_s_doclaynet/config.py \
# mmdetectionの保存したモデルcheck point
/workspaces/doclaynet_deploy/rtmdet_s_doclaynet/epoch_16.pth \
# inputのダミー
/workspaces/doclaynet_deploy/rtmdet_s_doclaynet/demo_input.png \
# onnxモデルの保存先
--work-dir mmdeploy_model/rtmdet
実行結果のSumally
05/18 16:28:25 - mmengine - INFO - Start pipeline mmdeploy.apis.pytorch2onnx.torch2onnx in subprocess
05/18 16:28:26 - mmengine - WARNING - Failed to search registry with scope "mmdet" in the "Codebases" registry tree. As a workaround, the current "Codebases" registry in "mmdeploy" is used to build instance. This may cause unexpected failure when running the built modules. Please check whether "mmdet" is a correct scope, or whether the registry is initialized.
05/18 16:28:26 - mmengine - WARNING - Failed to search registry with scope "mmdet" in the "mmdet_tasks" registry tree. As a workaround, the current "mmdet_tasks" registry in "mmdeploy" is used to build instance. This may cause unexpected failure when running the built modules. Please check whether "mmdet" is a correct scope, or whether the registry is initialized.
Loads checkpoint by local backend from path: /workspaces/doclaynet_deploy/rtmdet_s_doclaynet/epoch_16.pth
05/18 16:28:27 - mmengine - WARNING - DeprecationWarning: get_onnx_config will be deprecated in the future.
05/18 16:28:27 - mmengine - INFO - Export PyTorch model to ONNX: mmdeploy_model/rtmdet/end2end.onnx.
05/18 16:28:27 - mmengine - WARNING - Can not find torch._C._jit_pass_onnx_autograd_function_process, function rewrite will not be applied
~
05/18 16:28:31 - mmengine - INFO - Execute onnx optimize passes.
05/18 16:28:32 - mmengine - INFO - Finish pipeline mmdeploy.apis.pytorch2onnx.torch2onnx
05/18 16:28:32 - mmengine - INFO - Start pipeline mmdeploy.apis.utils.utils.to_backend in main process
05/18 16:28:32 - mmengine - INFO - atc --model=mmdeploy_model/rtmdet/end2end.onnx --framework=5 --output=mmdeploy_model/rtmdet/end2end --soc_version=Ascend310 --input_format=NCHW --input_shape=input:1,3,640,640
変換されたモデルを可視化してみると、実際計算グラフが見れるのがわかります。
実際にruntimeでうごくかどうかはonnxの対応次第なところがあるので、祈りましょう。
これでおしまいです。
onnxこわれた
出力されたonnxモデルの最後にくっついてるBatchMultiClassNMSという演算グラフが公式のONNXに存在せず( Ascend CANN という華為関係?のところに記述があるのみ)、このままでは動かないのでsne4onnxを使ってこの演算だけ消して対処しました。
使い方は書いてあるとおりで、次のようにインストールし、コマンドライン上で操作できます。
pip install onnx_graphsurgeon --index-url https://pypi.ngc.nvidia.com
pip install sne4onnx
sne4onnx \
--input_onnx_file_path rtmdet_s.onnx \
--input_op_names input \ # inputの名前
--output_op_names 1045 1017 \ # 作成するモデルのoutputにしたいグラフ上の入出力名
--output_onnx_file_path hitnet_sf_finalpass_720x960_head.onnx
これにより[1x8400x10]の(おそらく)クラス予測と[1x8400x4]のBBox予測が得られるモデルが完成したので、推論時にもともとついていた処理にあたる以下の設定のNMSを行ってあげれば良い。
iou_θ = 0.65
score_θ = 0.001
maxsize/cls = 200
maxtotalsize = 200
IR version関連で動いたり動かなかったりします。祈りが足りなかったか...
Discussion