【Detectron2を掘り下げる】val lossを計算・ロギングする

# はじめに
Detectron2は物体検出やセグメンテーションをする上で色々と便利な機能が簡単に実装できるライブラリです。
ただ、使ってみるとモデルを構築する上で提供されていないものがチラホラあるので、掘り下げて実装してみました。

今回は検証用（val）データセットに対してのlossを計算、ログを取る機能を実装してみました。
分類タスクではtrainとvalのデータセットに対してlossを計算して、一般的にはvalのlossの値をモニタリングして一番いいモデルを選択すると思います。
Detectron2ではtrainに対してのlossの計算はツールが提供されていますが、valに対してのツールは実装されておらず、自ら実装する必要があります。

# 環境

- Python 3.10.13
- CUDA 11.7
- pytorch==1.13.1
- detectron2==0.6

# 実装
### 1. configを設定する

```python
from detectron2.config import get_cfg
 
cfg = get_cfg()
cfg.DATASETS.TEST = ("your_validation_dataset",)
cfg.TEST.EVAL_PERIOD = 100
```

この例では100回のイテレーション後に1回、検証用データセットに対して評価が実施されます。

### 2. Evaluatorを実装する

今回はCOCOフォーマットのデータセットを使用しているため`COCOEvaluator`を実装します。
`DefaultTrainer`を継承した`Trainer`に以下のように`COCOEvaluator`を実装します。`TEST.EVAL_PERIOD`を設定すると、検証用データセット全体を使用してevaluatorが呼び出され、結果がストレージに書き込まれます。

```python
class Trainer(DefaultTrainer):

    @classmethod
    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
        if output_folder is None:
            output_folder = os.path.join(cfg.OUTPUT_DIR,"inference")
        return COCOEvaluator(dataset_name, cfg, True, output_folder)
```

### 3. val lossを計算するカスタムフックを実装する

```python
import datetime
import logging
import time
 
import detectron2.utils.comm as comm
import numpy as np
import torch
from detectron2.engine import HookBase
from detectron2.utils.logger import log_every_n_seconds
 
 
class LossEvalHook(HookBase):
    def __init__(self, eval_period, model, data_loader):
        self._model = model
        self._period = eval_period
        self._data_loader = data_loader

    def _do_loss_eval(self):
        total = len(self._data_loader)
        num_warmup = min(5, total - 1)

        start_time = time.perf_counter()
        total_compute_time = 0
        losses = []
        for idx, inputs in enumerate(self._data_loader):
            if idx == num_warmup:
                start_time = time.perf_counter()
                total_compute_time = 0
            start_compute_time = time.perf_counter()
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            total_compute_time += time.perf_counter() - start_compute_time
            iters_after_start = idx + 1 - num_warmup * int(idx >= num_warmup)
            seconds_per_img = total_compute_time / iters_after_start
            if idx >= num_warmup * 2 or seconds_per_img > 5:
                total_seconds_per_img = (
                    time.perf_counter() - start_time
                ) / iters_after_start
                eta = datetime.timedelta(
                    seconds=int(total_seconds_per_img * (total - idx - 1))
                )
                log_every_n_seconds(
                    logging.INFO,
                    "Loss on Validation  done {}/{}. {:.4f} s / img. ETA={}".format(
                        idx + 1, total, seconds_per_img, str(eta)
                    ),
                    n=5,
                )
            loss_batch = self._get_loss(inputs)
            losses.append(loss_batch)
        mean_loss = np.mean(losses)
        self.trainer.storage.put_scalar("val_loss", mean_loss)
        comm.synchronize()

        return losses

    def _get_loss(self, data):
        metrics_dict = self._model(data)
        metrics_dict = {
            k: v.detach().cpu().item() if isinstance(v, torch.Tensor) else float(v)
            for k, v in metrics_dict.items()
        }
        total_losses_reduced = sum(loss for loss in metrics_dict.values())
        return total_losses_reduced

    def after_step(self):
        next_iter = self.trainer.iter + 1
        is_final = next_iter == self.trainer.max_iter
        if is_final or (self._period > 0 and next_iter % self._period == 0):
            self._do_loss_eval()
        self.trainer.storage.put_scalars(timetest=12)
```

`val_loss`として出力されます。

### 4. Trainerにhookとして加える

`DefaultTrainer`を継承した`Trainer`に`LossEvalHook`をhookとして加えます。

```python
from detectron2.data import DatasetMapper, build_detection_test_loader
from detectron2.engine import DefaultTrainer
 
from src.hooks import LossEvalHook


class Trainer(DefaultTrainer):
    ...

    @classmethod
    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
        if output_folder is None:
            output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
        return COCOEvaluator(dataset_name, cfg, True, output_folder)

    def build_hooks(self):
        cfg = self.cfg.clone()
        hooks = super().build_hooks()
        # calculate and log val loss
        hooks.insert(
            -1,
            LossEvalHook(
                cfg.TEST.EVAL_PERIOD,
                self.model,
                build_detection_test_loader(
                    self.cfg, self.cfg.DATASETS.TEST[0], DatasetMapper(self.cfg, True)
                ),
            ),
        )
```

# 学習を実行！

学習が進み、`cfg.TEST.EVAL_PERIOD`で設定したイテレーションごとにvalデータセットに対して以下のようにval lossが計算されます。
※以下はMLflowでログを取った結果です。

モデルチェックポイントを実装すれば、val_lossが最小になったときのベストモデルを保存できます。

![スクリーンショット 2023-11-14 190844.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/471633/fcbc7b24-f717-bba3-e722-2e410ab8cadb.png)