YOLOXで多量物体検出訓練を行う設定

• １画像に数百単位（100以上）の物体があって、それをregressionやkeypointじゃなくて敢えてobject detection、特にyoloxで検出したいもの好きな人のための情報。

• 「デフォルトパラメーターで訓練しても全く検出できない。YOLOX使えないじゃん」ではなく、自分がちゃんと公式ドキュメントやissueを読んでいなかっただけなので、戒めとして公開しておく。

1. expファイル（Exp Class）の設定

python custom_yolo_exp.py

class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        # 他の設定は省略
        self.data_num_workers = 1 # or 2. default is 4
        self.enable_mixup = False
    
    def get_data_loader(self, batch_size, is_distributed, no_aug=False, cache_img: str = None):
        # 省略   
        self.dataset = MosaicDetection(
            # 省略
            preproc=TrainTransform(
                #max_labels=120, 
                max_labels=3000,
            # 省略
        )
        # 省略
        return train_loader

    def get_evaluator(self, batch_size, is_distributed, testdev=False, legacy=False):
        from yolox.evaluators import COCOEvaluator

        return COCOEvaluator(
            # dataloader=self.get_eval_loader(batch_size, is_distributed,
            #                                testdev=testdev, legacy=legacy),
            # force to 1 image per gpu
            dataloader=self.get_eval_loader(4, is_distributed,
                                            testdev=testdev, legacy=legacy),
            # 省略
        )

解説

• ラベル量（データ量）が多いと、データセット準備時（init prefetcher）にcore dumpする場合がある。その対策としてdata_num_workersを１に指定する。参考
• get_data_loaderメソッドをオーバーライドして、max_labelsの値を120から任意の値（１画像枚のラベル数の上限）に変更する。yolox/exp/yolox_base.pyよりget_data_loaderメソッドをコピペして該当箇所を修正すればよい。
• validation時のバッチサイズはtrain.pyで指定した訓練時の-b argumentと同一の値となる。validation時にメモリエラーになる場合、get_data_loaderと同様にyolox_base.pyよりget_evaluatorを引っ張ってきて、batchsizeを手動で最小値に指定する（複数gpuの場合にはgpu数x1などとする）。
• mixupはオフにする。ただでさえラベルが多いのにmixupで画像を重ね合わせるとラベルがさらに増えてしまう。（設定上許容できるならいじる必要ないけれども）

2. validation metrics計算時の設定

yoloxが使用するpycocotoolsでは１画像毎の最大ラベル数が100とされており、それ以上のラベルがあるデータセットで精度検証をすると見かけ上のARやAPが著しく低下する。例えば1画像500ラベル程度ある場合、APやARが0.3で頭打ちになる。これではbest_valのチェックポイントを拾ってくることができない。そのために下記coco_evaluator.pyに手を加える。

yolox/evaluators/coco_evaluator.py

try:
    from yolox.layers import COCOeval_opt as COCOeval
except ImportError:
    from pycocotools.cocoeval import COCOeval

    logger.warning("Use standard COCOeval.")

# ここにmaxDetsを変更した新規継承クラスを定義し、cocoEvalにはそのクラスを用いるように記述する。
logger.info("overide COCOeval")
class CustomCOCOeval(COCOeval):
    def __init__(self, cocoGt=None, cocoDt=None, iouType='segm'):
        super().__init__(cocoGt, cocoDt, iouType)
        self.params.maxDets = [100, 500, 1000]
# cocoEval = COCOeval(cocoGt, cocoDt, annType[1])
cocoEval = CustomCOCOeval(cocoGt, cocoDt, annType[1])

解説

• pycocotools（yoloxレポジトリのCOCOeval_opt含む）はデフォルトで１画像毎のラベル数を[1,10,100]とハードコードしていて、APやARの計算には100ラベル時の値が用いられる。ラベル数はCOCOEval.params.maxDetsで管理されているので、ここを改変すれば良い。COCOeval_optもCOCOevalを継承しているので同様に変更可能。ここでは[1,10,100] -> [100,500,1000]とした。1000固定でもよいがmaxDets[2]などとリスト処理前提のコードが複数箇所あるので３要素リストの形式を保つ。

結果

2023-08-15 11:00:46 | INFO     | yolox.core.trainer:354 - 
Average forward time: 27.81 ms, Average NMS time: 2.06 ms, Average inference time: 29.88 ms
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.296
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=1000 ] = 0.947
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=1000 ] = 0.915
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=1000 ] = 0.322
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=1000 ] = 0.799
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=1000 ] = 0.836
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.305
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=500 ] = 0.798
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=1000 ] = 0.798
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=1000 ] = 0.334
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=1000 ] = 0.839
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=1000 ] = 0.902
per class AP:
| class   | AP     |
|:--------|:-------|
| seed    | 75.995 |
per class AR:
| class   | AR     |
|:--------|:-------|
| seed    | 79.796 |

以上

3. （おまけ）multi gpuでvalidation phase終了時にフリーズする場合

pip install -v -e .