【機械学習】EfficientDet(effdet)を使った物体検知・分類の実装

前回の記事では、SIIM Covid-19 detectionコンペのデータを使ってEfficientDetと呼ばれる物体検知・分類モデルの学習を実装しようしているということについて書いたわけだが、とりあえずそれらしい値を返すものができたので、本日はeffdetに関連する部分で注意すべき点をメモしておこうと思う。

2023年4月時点で動作したコードはこちらである。なお当方、開発のプロではないため、非効率な部分も多々あると思う。何かお気づきの際はコメントでお知らせいただけるとありがたいです。

1. ライブラリのインストール
2. Datasetの定義
3. collate_fnの定義
4. effdetモデルの設定と学習

1. ライブラリのインストール

まずはライブラリのインストールだが、注意すべきはオリジナルのeffdetライブラリを使おうとすると、何か所かでエラーが発生することだ。

1つはコスト関数の計算時にエラーが発生する。詳細はこちら（view size is not compatible with input tensor's size and stride）を参照いただきたいが、テンソルのreshapeを行うところでviewを使っているため、reshapeを使えとのことだ。

もう1つはこちら(effdet内、anchors.pyで生じたエラー)である。1画像内に複数バウンディングボックスが含まれる場合に発生するため修正を行った。（ただし、この記事にある修正を行うと、GPUを使って学習する際に別のエラーが生じるので、別の修正方法を考える必要がある。）

Kaggleのプラットフォーム上ではインストールしたライブラリ内のファイルを編集するのが難しいため、あらかじめ編集したものを'effdet-030-package-dataset'という名前のデータセット上にアップロードして、そちらから読み込むようにしている。

なお、effdetを使うにはomegaconf, timm, huggingface_hub, antlr4, pycocotoolsというライブラリも必要だが、omegaconf, timm, huggingface_hubはアップロードしたものを読み込んでおり、その他はpipでインストールするようにしてある。

!pip install antlr4-python3-runtime==4.9.3
!pip install pycocotools==2.0.2
!pip install /kaggle/input/effdet-030-package-dataset/packages/huggingface_hub-0.13.3-py3-none-any.whl
import sys
sys.path.insert(0, "../input/effdet-030-package-dataset/packages/")

2. Datasetの定義

学習データをtorch.utils.data.Datasetの子クラスとして定義する。ここでは次の2点がポイントとなる。

Datasetが画像、正解バウンディングボックス、正解ラベルを返すように設定する

Datasetクラスの__getitem__()関数を定義して、学習データの画像、正解バウンディングボックス、正解ラベルを返すように設定する。いずれもtorch.Tensor型として返すようにするわけだが、画像は単体で、バウンディングボックスとラベルはdict型としてまとめて返すように設定し、バウンディングボックスは'bbox'、ラベルは'cls'というキーに対応して返すようにする。画像内にバウンディングボックスが1つの場合、bboxには(1, 4)サイズ、clsには（1）サイズのテンソルを格納する。今回は画像内に複数のバウンディングボックスが含まれるため、bboxには(n, 4)サイズ、clsには(n, 1)サイズ（nはバウンディングボックスの数）のテンソルを格納する。

### __getitem__()関数内の処理
# 画像は変数imageに(3, width, height)サイズのnumpy.arrayとして格納されている
# バウンディングボックスは変数boxesに(n, 4)のnumpy.arrayとして格納されている
# ラベルは変数labelsに(n, 1)のnumpy.arrayとして格納されている

import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2

# 画像をリサイズしてtorch.Tensor化する処理の定義、バウンディングボックスも併せてリサイズされるように設定
albu = A.Compose([
    A.Resize(width=self.image_size, height=self.image_size, p=1),
    ToTensorV2()
], bbox_params=A.BboxParams(format='pascal_voc', label_fields=['labels']))

result = albu(
    image = np.array(image),
    bboxes = bboxes,
    labels = labels
)

# 座標の入れ替え from xyxy to yxyx
box = np.array(result['bboxes'])[:,[1,0,3,2]]

# バウンディングボックスとラベルのtorch.Tensor化
bboxes = torch.FloatTensor(box)
labels = torch.FloatTensor(np.array(result['labels']))

# バウンディングボックスとラベルをdictにまとめる
y = {
    'bbox': bboxes,
    'cls': labels
}
    
return x, y

バウンディングボックスの座標をyxyx形式に入れ替える

こちらの記事でも紹介したが、effdetにバウンディングボックスを入力する際には座標をyxyx形式に入れ替えてから入力する必要がある。これを忘れると学習後の予測結果が全く異なるものになるので注意が必要。一方で、画像のリサイズ時にバウンディングボックスも併せて座標の縮小を行うが、これを行うAlbumentationsはxyxy形式が必要。そのため、画像をリサイズし、Datasetからバウンディングボックスを出力する直前で座標を入れ替える必要がある。

3. collate_fnの定義

今回の学習データにおいては、1画像が複数のバウンディングボックスをもち、かつ画像ごとにバウンディングボックスの数が異なる。このような場合、複数データをバッチ化してモデルに入力する際、同じバッチ内に含まれるデータに対して、バウンディングボックスの数をそろえるように求められる。(参考：1枚の画像に対して複数のバウンディングボックスをもつ場合の処理)

このような場合、「2. Datasetの定義」で定義したDatasetを後でtorch.utils.data.DataLoaderを使ってバッチ化するのだが、その際にバウンディングボックスの数をそろえる処理を記述すればよい。具体的には、バウンディングボックスが少ない画像に対して、すべての座標が0、ラベルが0のバウンディングボックスを追加することで数を合わせる。

def pad_collate_fn(batch):
    # batch内のデータの中で、バウンディングボックスの最大数を計算
    shapes = [item[1]['bbox'].shape[0] for item in batch]
    max_shape = max(shapes)
    
    # 足りないバウンディングボックスを0パディングする
    # xは画像データ（今は使わない）
    # yはdictで、バウンディングボックス（キー：bbox）とラベル（キー：cls）を含む
    padded_batch = []
    for x, y in batch:
        if any(elem == 0 for elem in y['cls']):
            continue
        pad_size = max_shape - y['bbox'].shape[0]
        bbox_padding = [0, 0, 0, pad_size]
        cls_padding = [0, 0, 0, pad_size]
        padded_y = {
            'bbox': F.pad(y['bbox'], bbox_padding, mode='constant', value=0),
            'cls': F.pad(y['cls'].reshape((y['cls'].shape[0],1)), cls_padding, mode='constant', value=0)
        }
        padded_batch.append((x, padded_y))
    
    return default_collate(padded_batch)

# datasetは定義された学習データを呼び出すためのtorch.utils.data.Datasetの子クラス
loader = DataLoader(
    dataset, batch_size=4, num_workers=2, collate_fn=pad_collate_fn
)

4. effdetモデルの設定と学習

effdetモデルを設定するには、まずget_efficientdet_configを使って学習済パラメータを読み込み、EfficientDetでモデルを設定する。学習の際のフォワードプロパゲーションはDetBenchTrainを使って実行する。分類したいクラスの数は読み込んだパラメータのうちnum_classesというパラメータを変更することで設定可能。

# 学習済パラメータの読み込み
cfg = get_efficientdet_config(f'tf_efficientdet_d0')
# 分類したいクラスの数を変更
cfg.num_classes = 3
# モデルの設定
model = EfficientDet(cfg)
# フォワードプロパゲーション用の関数
bench = DetBenchTrain(model)

# 学習ループ
from tqdm import tqdm
from torch import optim

optimizer = optim.Adam(model.parameters(), lr=1e-4)

for epoch in range(1, number_of_epochs):
    t = tqdm(loader, leave=False)

    for inputs, targets in t:
        optimizer.zero_grad()
        losses = bench(inputs, targets)
        loss = losses['loss']
        loss.backward()
        optimizer.step()
        
        t.set_description(f'{header}{message}')
        t.refresh()

ちなみに、学習後に予測値を得るにはDetBenchPredictを使って次のようにする。

# image: バウンディングボックスとその分類を検知したい画像、(3, width, height)サイズのnumpy.arrayとする

image = image.unsqueeze(0)

bench = DetBenchPredict(model)
with torch.no_grad():
    output = bench(image)

変数outputは（n, 6）のテンソルとなっている。（参照：学習後の予測精度）nは予測されたバウンディングボックスの数で、各列はそれぞれ

• 0-3列目：バウンディングボックスの座標、xyxy形式
• 4列目：0～1の値をとり、予測の精度を表す
• 5列目：予測されたクラス

となっている。

さて、残る問題は得られた結果をどのようにして「予測結果」としてまとめるかである。ひとまず本日はここまで。