More than 5 years have passed since last update.

PyTorchで始める物体検出:Yolo 9000 Better, Faster, Stronger

Last updated at 2018-10-24Posted at 2017-11-14

誰向け

高速に物体検出をしたい人
論文の内容を理解したい人
コードで実際に動作させてみたい人
コードと論文の関係性を把握したい人

先に読んだ方が良い記事

物体検出に関する用語が多いためDeepに理解する深層学習による物体検出 by Kerasのssdの説明より前の部分を読むと理解の準備が出来ると思います。

Yolo9000が達成したことを3行で表すなら

精度向上：バッチ正規化、アンカーの追加、高解像度、シグマ関数による制約、K-meanを用いた最適なアンカーボックスの決定、フィルターの増加、マルチスケーリングなどによって精度向上
高速化：パラメータが少なくなるDarkNetを採用し計算量を減らし、さらに事前学習モデルによって学習時間を短くした
分類クラスの増加(9000)：検出用のデータセットがなくても同一のドメインのデータセットがあれば多数のクラスの物体検出が可能なことを示した。この論文ではImageNetとCOCOを組み合わせて学習する方法を提案することで多数のクラスの物体検出を可能にした

論文は下記

学習済みモデルとコードは下記

Yolo 9000とは

物体検出において高速で精度が高いものが求められていますが、現状では小さなデータ・セット（クラス数21程度など）でのみしか確認されていません。本来は数千の単位のデータを識別したいのですが学習データの作成コストは少くないです。
この手法の利点は階層型の物体識別を用いることによって現状の物体検出を拡張可能にしたこと、検出と識別を同時最適化を行うことです。

リアルタイムで9000カテゴリ以上の物体をリアルタイムで検出し、2016.12.25でState of Artsを達成しています。（すこし古い）

Better

Fast-RCNNに比べ、Yoloでは検出位置のRecallが低かったのでそれを改善するようにフォーカスしています。
近年では深く、大きいネットワークの方が精度がよくなる傾向にありますがYolov2は精度と速度を保つためにあえてそのようなアプローチはしていません。

精度向上の工夫

Batch Normalization: すべての畳込み層にBatch NormalizationをいれてmAPを2％向上させています。

PyTorch上では下記の部分で実装されています。

class Conv2d_BatchNorm(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, relu=True, same_padding=False):
        super(Conv2d_BatchNorm, self).__init__()
        padding = int((kernel_size - 1) / 2) if same_padding else 0

        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=padding, bias=False)
        self.bn = nn.BatchNorm2d(out_channels, momentum=0.01)
        self.relu = nn.LeakyReLU(0.1, inplace=True) if relu else None

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        if self.relu is not None:
            x = self.relu(x)
        return x

高解像度識別器：Yoloでは224の x 224の画像サイズをYolo v2では448 x 448にしてImageNetを事前学習モデルとして使用することでmAPで4％精度向上しています。
実際に使っているのは448x448ではなく416x416でこれは32倍のダウンサンプリングをして１３x13にするのに都合の良い数字であること、奇数を最終層にすることによって中心位置を決められる利点から変えています。448を３２倍のダウンサンプリングすると14で偶数となり中心位置が決められません

中心位置の利点：大きいオブジェクトの場合は中心位置の近くに位置する割合が高いので後述するアンカーボックスを決定するのが容易になります。

inp_size = np.array([416, 416], dtype=np.int) # w, h

畳み込みとアンカーボックス：Yoloでは畳込み層の特徴量を全結合して使用していましたが、アンカーボックスを使用することによってアンカーボックスごとにIoUを計算するようにして学習を容易にできます。Yoloでは98のボックスで済みましたがAnchorボックスを使用すると1000程度になるためmapは0.3%程度減少しますがRecallは7％ほど上昇しています。

python上でのアンカーの定義（論文ではIoUを用いたK-meansにより最適なものを選んでいるため下記が最適なのかは不明）

anchors = np.asarray([(1.08, 1.19), (3.42, 4.41), (6.63, 11.38), (9.42, 5.11), (16.62, 10.52)], dtype=np.float)
num_anchors = len(anchors)

2x2の画像を用意してx,y,h,wのオフセットに対してアンカーボックスを適用した場合は下記のようにデータセットを用意します。

hw = 4
_boxes = np.zeros([hw, num_anchors, 4], dtype=np.float)

上記で出力されるデータは下記のようになります

array([
       # この部分は各画像の座標、縦軸がアンカーボックスの数、横軸がオフセットの座標
       [[ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.]],

       [[ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.]],

       [[ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.]],

       [[ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.]]])

Dimension Clusters：アンカーボックスをIoUスコアを用いたk-meansによりクラスタリングし、クラスターの平均のIoUスコアを算出し、クラスごとに優先順位をつけて学習すると手動で9クラス優先順位付けしたものよりも5クラスでクラスタリングした方が0.1%ほどIoUが高かったようです。
下記の図はクラスター数と精度の関係性を表します。コードではすでにアンカーボックスのサイズが決められていました。

Direct location prediction：位置の予測に制約がないため学習が安定するのに時間がかかるためシグマ関数で制約をかけています。

p_w, p_hはアンカーボックスのサイズ \\
b_x = \sigma(t_x) + c_x \\
b_y = \sigma(t_y) + c_y \\
b_w = p_we^{(t_w)} \\ 
b_h = p_he^{(t_h)} \\
Pr(object) ∗ IOU(b, object) = \sigma(t_o)

実際のコードでは下記のように記述しています。

        xy_pred = F.sigmoid(conv5_reshaped[:, :, :, 0:2])
        wh_pred = torch.exp(conv5_reshaped[:, :, :, 2:4])

Fine-Grained Features.： 13 x 13の特徴マップは大きな画像からの小さな物体を検出するには不向きなため、もっと高解像度な情報が必要です。ResNetと異なり、Yoloでは26 x 26 x 512の特徴マップを13 x 13 x 2048にしてオリジナルの特徴量に合体しています。これによって1％の精度向上をしています。

下記がネットワークの設定です。

        net_cfgs = [
            # conv1s
            [(32, 3)],
            ['M', (64, 3)],
            ['M', (128, 3), (64, 1), (128, 3)],
            ['M', (256, 3), (128, 1), (256, 3)],
            ['M', (512, 3), (256, 1), (512, 3), (256, 1), (512, 3)],
            # conv2
            ['M', (1024, 3), (512, 1), (1024, 3), (512, 1), (1024, 3)],
            # ------------
            # conv3
            [(1024, 3), (1024, 3)],
            # conv4
            [(1024, 3)]
        ]

conv1は最終層のフィルター数が512のためその特徴量を用います。
この部分で特徴量を合わせています。


self.conv4, c4 = _make_layers((c1*(stride*stride) + c3), net_cfgs[7])

実際にforward処理で動作している部分は下記です。reorgはフィルター数を指定の数だけ倍数化する変わリにサイズをその倍数だけ小さくします。この処理でconv1sとconv3のサイズを同一にしてconcatが可能になります。

        conv1s_reorg = self.reorg(conv1s)
        cat_1_3 = torch.cat([conv1s_reorg, conv3], 1)

self.reorgの処理部分が下記でフィルターサイズを4倍にしています。


 out_w, out_h, out_c = int(w / stride), int(h / stride), c * (stride * stride)

Multi-ScaleTraining.: 今回は416x416の画像サイズですがマルチスケールで学習した方がロバスト性が向上するのでYolov2では10バッチごとに32の倍数の画像イメージサイズをランダムに選び学習している{320,・・・608}までの候補で学習しています。
これは畳み込みとプーリング層のみで構成されているため可能です。全結合層があると特定のサイズに入力を固定しないと出来ないですがないのでこのような処理が可能になります。
下記の図がFPSとmaPの関係性を表しています。

コード上ではマルチスケーリングの実装は確認出来ませんでしたがDefine by RunのフレームワークであるPyTorchであれば実装は簡単に行うことが可能です。

実験結果

下記の図波既存の他のモデルとの比較になります。Yolov2が高速かつ高精度ということが分かります。

下記の図は各処理をかけた場合の効果を表しています。

Faster

自動運転やロボットなどにおいて速度が重要です。

VGG-16に精度2%ほど劣るがGoogle Netのアーキテクチャをベースにカスタマイズしたモデルで306.9億のfloat操作を85.2億の操作まで下げることに成功しました。

Darknet-19

VGG-modelと同様に3x3のフィルター、各プーリング処理の後にチャネル数を２倍にし、グローバル平均プールを使用して予測部分を圧縮します。また学習の安定、正則化のためにバッチノーマライザーションを使用します。
55.8億操作で済みます。(Google Netの6分の1の操作)

コード上でのDarkNetの定義部分

Training for classification

224x224のサイズでImageNet 1000クラス分類データセットでData augmentationありで160epochで学習しそのあとに解像度を上げて448のサイズで学習することでわずか10epochで学習可能。（学習率は0.001)

学習済みモデルの読み込み部分

net.load_from_npz(cfg.pretrained_model, num_conv=18)

重みの読み込み部分
畳み込みの部分に重みを移す処理をしています。

Training for detection.

Yoloのモデル

最終層のレイヤー（3 x 3の1024フィルターの畳み込みレイヤーと全結合相）を3 x 3の512フィルターの畳み込みレイヤーと3 x 3の1024フィルターの畳み込みレイヤーの組み合わせに入れ替えています。
最終層は検出のために1 x 1のカーネルサイズを設定して出力はアンカーの数 * (クラス数 + 5)
上記の5はオフセットとIoUの値

検出用のコードは下記

out_channels = cfg.num_anchors * (cfg.num_classes + 5)
self.conv5 = net_utils.Conv2d(c4, out_channels, 1, 1, relu=False)

Yoloの場合は全結合層があるため計算量が多くなるがYolov2では排除しているため計算量が減り高速になっています。

全結合層があることによるパラメータ量(一つ前の畳み込み層*全結合層)

(3 x 3 x 1024) x 4096 = 37,748,736(約4000万)

Yolov2の場合は(一つ前の畳み込み層*フィルター数)

(3 x 3 x 1000) x 1000 = 9,000,000(約900万)

本当に減っているかはトータルのパラメータ数を計算する必要がありますが全結合層をなくすことにより4倍程度の削減効果があります。

計算例は下記の72ページを参考にしています。

Stronger

この論文で述べているStrongerはより多くのクラスを識別できることです。
現状のアーキテクチャではクラス間が排他的なため犬の中でも「ノーフォーク・テリア」、「ヨークシャー」などの犬種は異なるが犬のクラスを排他的な枠組みで導出してしまうことになります。
この章ではどのようにそれを避けるかについて述べています。

Hierarchical classification

WordNetの構造とほぼ同一でImageNetは複雑な文字情報を扱うために有効グラフで表されています。WordTreeと呼ばれる階層構造の木を作成するためImageNetで使用されている視覚名詞からルートノード（今回は物理オブジェクト）までのパスを出します。
WordNedで使用されているノードはsynset(類義語のひとかたまり)であり、1つのパスのみが大半なのでそれを階層的な木に割り当てて行き、その操作を再起的に行います。複数のパスがある場合は短いパスを選択します。
この操作によってWordTreeを作成します。

WordTree

WordTreeでのクラス識別は条件付き確率の計算で導出します。Pr(animal\ |\ physical object)は1です。

Pr(Norfolk terrier) = Pr(Norfolk terrier\ |\ terrier) ∗Pr(terrier\ |\ hunting dog) ∗ . . .∗
Pr(mammal\ |\ animal) ∗Pr(animal\ |\ physical object)

下記の図だとImageNetは各クラスをsoftmaxで計算しているのに対して、WodTreeは階層ごとにsoftmaxを適用していることが分かります。

この手法で71.9%の精度を達成（Top 1)、上位5件では90.4%と1000クラス分類ではかなりの高精度を出しています。

実装には反映されていないので実際に実装する場合は下記のようになると思います。

まだ試せていないですが（汗

softmax_word_tree = []
for wordnet_index in output_index:
    softmax_word_tree.append(F.softmax(output[5:wordnet_index]))
output_tensor = torch.cat(tuple(softmax_word_tree), 0)

Dataset combination with WordTree.

WordTreeを使用することでCOCOとImageNetの異なるデータセットも扱えるようになります。WordTreeによってどの位置に各データのオブジェクトがあるかを把握できます。

Joint classification and detection.

ImageNetの多様なクラス情報をクラス識別の学習に使用し、COCOの位置情報を物体検出に使用しています。

WordTreeを使用することでCOCOとImageNetの学習データが使用可能になります。ImageNetの方が大きなモデルなのでCOCOは4倍のオーバーサンプリングをして学習をしています。
ラベルの情報がなくラベルが一致しないものはエラーとして学習します。

クラスロス計算は誤差逆伝搬は0.3IOU検出できている過程の元でクラス誤差のみ伝搬させます。

COCOとImageNetでお互いに一致していないデータ156オブジェクトに対して19.7 mAPを達成しています。

動物ではCOCOのデータがあるので上手く動作しますが衣服や装備では上手く行きませんでした。典型例がサングラスと水泳のゴーグルです。ドメインが異なる点からだと思われます。

実装について

Pytorchとは

特徴としては簡潔にかくとDefine by Runがもっとも良い利点です。

データによってモデルの構造を変更したい（簡潔なデータは浅いモデル、難しいデータは深いモデル）の場合に有効です。
自然言語処理では文の長さが変わるので需要がありそうです。
画像であれば入力画像のサイズを変えて学習することも可能です。

この手法によって画像サイズが大きい、長い文のようなデータは層の深い表現力が大きいモデル、小さいサイズ、短い文は層の小さい表現力の小さなモデルにすることが可能です。
今回のYoloではマルチスケールのモデルを作成していたのでPyTorchのDefine by Runは効果的です。下記がmnistで画像のサイズを28と56に変更した場合のネットワークです。


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)
        self.conv1_2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_2 = nn.Conv2d(20, 40, kernel_size=5)
        self.fc1_2 = nn.Linear(360, 50)
        self.over_size = 28

    def forward(self, x):
        _, _, h, w = x.size()
        if h > self.over_size:
            x = F.relu(F.max_pool2d(self.conv1(x), 2))
            x = F.relu(F.max_pool2d(self.conv1_2(x), 2))
            x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2_2(x)), 2))
            x = x.view(-1, 360)
            x = F.relu(self.fc1_2(x))
        else:
            x = F.relu(F.max_pool2d(self.conv1(x), 2))
            x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
            x = x.view(-1, 320)
            x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x)

上記の例では若干ですが層を深くしています。Define and Runだとこのような処理はできません。

実際に試したコードは下記です。

学習の全体像

データローダーの設定

下記で設定しています。学習中にnext_batchで処理を呼び出していますがImageDatasetを継承しているためnext_batchの処理はImageDatasetに記述されています。

imdb = VOCDataset(cfg.imdb_train, cfg.DATA_DIR, cfg.train_batch_size,
                  yolo_utils.preprocess_train, processes=2, shuffle=True, dst_size=cfg.inp_size)

next_batchは下記です。メモリーの効率化のためgeneratorによるバッチサイズごとのアクセス、並列処理によってすべてのデータセットに同一の処理をかけることによる高速化をしています。

    def next_batch(self):
        batch = {'images': [], 'gt_boxes': [], 'gt_classes': [], 'dontcare': [], 'origin_im': []}
        i = 0
        while i < self.batch_size:
            try:
                images, gt_boxes, classes, dontcare, origin_im = self.gen.next()
                batch['images'].append(images)
                batch['gt_boxes'].append(gt_boxes)
                batch['gt_classes'].append(classes)
                batch['dontcare'].append(dontcare)
                batch['origin_im'].append(origin_im)
                i += 1
            except (StopIteration, AttributeError):
                indexes = np.arange(len(self.image_names), dtype=np.int)
                if self._shuffle:
                    np.random.shuffle(indexes)
                self.gen = self.pool.imap(self._im_processor,
                                          ([self.image_names[i], self.get_annotation(i), self.dst_size] for i in indexes),
                                          chunksize=self.batch_size)
                self._epoch += 1
                print('epoch {} start...'.format(self._epoch))
        batch['images'] = np.asarray(batch['images'])

        return batch

ネットワークの設定

net = Darknet19()

同一の処理が多いためネットワークの層の構築は下記のメソッドにより効率化しています。

def _make_layers(in_channels, net_cfg):
    layers = []

    if len(net_cfg) > 0 and isinstance(net_cfg[0], list):
        for sub_cfg in net_cfg:
            layer, in_channels = _make_layers(in_channels, sub_cfg)
            layers.append(layer)
    else:
        for item in net_cfg:
            if item == 'M':
                layers.append(nn.MaxPool2d(kernel_size=2, stride=2))
            else:
                out_channels, ksize = item
                layers.append(net_utils.Conv2d_BatchNorm(in_channels, out_channels, ksize, same_padding=True))
                # layers.append(net_utils.Conv2d(in_channels, out_channels, ksize, same_padding=True))
                in_channels = out_channels

    return nn.Sequential(*layers), in_channels

学習済みモデルの読み込み

学習済みモデルの読み込みをして畳み込み層の重みを下記で取得しています。

net.load_from_npz(cfg.pretrained_model, num_conv=18)

    def load_from_npz(self, fname, num_conv=None):
        # 重みを指定
        dest_src = {'conv.weight': 'kernel', 'conv.bias': 'biases',
                    'bn.weight': 'gamma', 'bn.bias': 'biases',
                    'bn.running_mean': 'moving_mean', 'bn.running_var': 'moving_variance'}
        # 学習済みモデルの読み込み
        params = np.load(fname)
        # モジュールのすべての状態をdictionary形式で返却
        #     http://pytorch.org/docs/master/nn.html?highlight=state_dict#torch.nn.Module.state_dict
        own_dict = self.state_dict()
        # モデルのキー(conv_weightなど)を取得
        keys = list(own_dict.keys())

        # 畳み込み層を効率よくアクセスするため5つ刻み
        for i, start in enumerate(range(0, len(keys), 5)):
            if num_conv is not None and i >= num_conv:
                break
            end = min(start+5, len(keys))
            for key in keys[start:end]:
                list_key = key.split('.')
                ptype = dest_src['{}.{}'.format(list_key[-2], list_key[-1])]
                src_key = '{}-convolutional/{}:0'.format(i, ptype)
                print((src_key, own_dict[key].size(), params[src_key].shape))
                param = torch.from_numpy(params[src_key])
　　　　　　　　　 # kernelのみ重みの配列の順序を変更
                if ptype == 'kernel':
                    param = param.permute(3, 2, 0, 1)
                own_dict[key].copy_(param)

実際の取得されるパラメータの様子

('0-convolutional/kernel:0', torch.Size([32, 3, 3, 3]), (3, 3, 3, 32))
('0-convolutional/gamma:0', torch.Size([32]), (32,))
('0-convolutional/biases:0', torch.Size([32]), (32,))
('0-convolutional/moving_mean:0', torch.Size([32]), (32,))
('0-convolutional/moving_variance:0', torch.Size([32]), (32,))
('1-convolutional/kernel:0', torch.Size([64, 32, 3, 3]), (3, 3, 32, 64))
('1-convolutional/gamma:0', torch.Size([64]), (64,))
('1-convolutional/biases:0', torch.Size([64]), (64,))
('1-convolutional/moving_mean:0', torch.Size([64]), (64,))
('1-convolutional/moving_variance:0', torch.Size([64]), (64,))
:

Forward処理

Forward処理は下記になります。

net(im_data, gt_boxes, gt_classes, dontcare)

ロスの計算のため教師データ、出力結果の変換が必要になります。
下記で予測したボックスに対してアンカーボックスを適用しています。

    anchors = np.ascontiguousarray(cfg.anchors, dtype=np.float)
    bbox_pred_np = np.expand_dims(bbox_pred_np, 0)
    bbox_np = yolo_to_bbox(
        np.ascontiguousarray(bbox_pred_np, dtype=np.float),
        anchors,
        H, W)
    bbox_np = bbox_np[0]  # bbox_np = (hw, num_anchors, (x1, y1, x2, y2))   range: 0 ~ 1
    bbox_np[:, :, 0::2] *= float(inp_size[0])  # rescale x
    bbox_np[:, :, 1::2] *= float(inp_size[1])  # rescale y

下記が出力サイズへの変更、アンカーボックスによるスケールの変更の処理です。

    for b in range(bsize):
        for row in range(H):
            for col in range(W):
                ind = row * W + col
                for a in range(num_anchors):
                    # 出力画像のサイズと合わせるため中央位置を計算
                    cx = (bbox_pred[b, ind, a, 0] + col) / W
                    cy = (bbox_pred[b, ind, a, 1] + row) / H
                    # 出力画像のサイズと合わせるため幅、高さを計算。0.5倍は中央位置からの幅、高さのため
                    bw = bbox_pred[b, ind, a, 2] * anchors[a][0] / W * 0.5
                    bh = bbox_pred[b, ind, a, 3] * anchors[a][1] / H * 0.5

                    # オフセットを計算。(x_min, y_min, x_max, y_max)
                    bbox_out[b, ind, a, 0] = cx - bw
                    bbox_out[b, ind, a, 1] = cy - bh
                    bbox_out[b, ind, a, 2] = cx + bw
                    bbox_out[b, ind, a, 3] = cy + bh

下記が予測したボックスと真のボックスの一致度合いを計測している部分です。

    bbox_np_b = np.reshape(bbox_np, [-1, 4])
    ious = bbox_ious(
        np.ascontiguousarray(bbox_np_b, dtype=np.float),
        np.ascontiguousarray(gt_boxes_b, dtype=np.float)
    )
    best_ious = np.max(ious, axis=1).reshape(_iou_mask.shape)
    iou_penalty = 0 - iou_pred_np[best_ious < cfg.iou_thresh]
    _iou_mask[best_ious <= cfg.iou_thresh] = cfg.noobject_scale * iou_penalty

一致度合いをIoUにより計算している部分です。

    for k in range(K):
        # query_boxesは真のボックスのエリア。下記は面積を計算している
        qbox_area = (
            (query_boxes[k, 2] - query_boxes[k, 0] + 1) *
            (query_boxes[k, 3] - query_boxes[k, 1] + 1)
        )
        for n in range(N):
            # 下記は真のボックスと予測したボックスの幅の一致している部分を導出している
            iw = (
                min(boxes[n, 2], query_boxes[k, 2]) -
                max(boxes[n, 0], query_boxes[k, 0]) + 1
            )
            if iw > 0:
                # 下記は真のボックスと予測したボックスの高さの一致している部分を導出している
                ih = (
                    min(boxes[n, 3], query_boxes[k, 3]) -
                    max(boxes[n, 1], query_boxes[k, 1]) + 1
                )
                if ih > 0:
                    # 一致している部分が幅も高さも1以上の場合は予測したボックスの面積を計算
                    box_area = (
                        (boxes[n, 2] - boxes[n, 0] + 1) *
                        (boxes[n, 3] - boxes[n, 1] + 1)
                    )
                    # 一致部分の面積を計算
                    inter_area = iw * ih
                    # 下記で一致率を計算
                    intersec[n, k] = inter_area / (qbox_area + box_area - inter_area)

下記で真のボックスのサイズを出力のサイズに変更しています。

    cell_w = float(inp_size[0]) / W
    cell_h = float(inp_size[1]) / H
    cx = (gt_boxes_b[:, 0] + gt_boxes_b[:, 2]) * 0.5 / cell_w
    cy = (gt_boxes_b[:, 1] + gt_boxes_b[:, 3]) * 0.5 / cell_h
    cell_inds = np.floor(cy) * W + np.floor(cx)
    cell_inds = cell_inds.astype(np.int)

    target_boxes = np.empty(gt_boxes_b.shape, dtype=np.float)
    target_boxes[:, 0] = cx - np.floor(cx)  # cx
    target_boxes[:, 1] = cy - np.floor(cy)  # cy
    target_boxes[:, 2] = (gt_boxes_b[:, 2] - gt_boxes_b[:, 0]) / inp_size[0] * out_size[0]  # tw
    target_boxes[:, 3] = (gt_boxes_b[:, 3] - gt_boxes_b[:, 1]) / inp_size[1] * out_size[1]  # th

下記で最も良いアンカーボックスを選択しています。

    gt_boxes_resize = np.copy(gt_boxes_b)
    # 真のボックスを出力のサイズに変換
    gt_boxes_resize[:, 0::2] *= (out_size[0] / float(inp_size[0]))
    gt_boxes_resize[:, 1::2] *= (out_size[1] / float(inp_size[1]))
    # 各アンカーボックスにおけるIoUを計算
    anchor_ious = anchor_intersections(
        anchors,
        np.ascontiguousarray(gt_boxes_resize, dtype=np.float)
    )
    # IoUが最も高いアンカーボックスを選択
    anchor_inds = np.argmax(anchor_ious, axis=0)

下記が具体的なアンカーボックスの選択のために各アンカーボックスにおけるIoUを計算する処理になります。

    for n in range(N):
        anchor_area = anchors[n, 0] * anchors[n, 1]
        for k in range(K):
            # 真のボックスの幅を計算
            boxw = (query_boxes[k, 2] - query_boxes[k, 0] + 1)
            # 真のボックスの高さを計算
            boxh = (query_boxes[k, 3] - query_boxes[k, 1] + 1)
            # 一致している幅を計算
            iw = min(anchors[n, 0], boxw)
            # 一致している高さを計算
            ih = min(anchors[n, 1], boxh)
            # 一致している面積を計算
            inter_area = iw * ih
            # 一致率を計算
            intersec[n, k] = inter_area / (anchor_area + boxw * boxh - inter_area)

下記で選択した最もIoUが高いアンカーボックスを適用して_boxes, _ious, _classes, _box_mask, _iou_mask, _class_maskを計算します。maskはスケーリング処理を行ってロス計算時の重要性を決定しています。

    for i, cell_ind in enumerate(cell_inds):
        # ネットワークの予測値を超える場合は下記の処理をしない
        if cell_ind >= hw or cell_ind < 0:
            print cell_ind
            continue
        # 最も良いアンカーボックスを選択
        a = anchor_inds[i]

        iou_pred_cell_anchor = iou_pred_np[cell_ind, a, :]  # 0 ~ 1, should be close to 1
        # マスク処理によってスケーリングを行い、ロス計算時の重要度を設定。一致率が高いほどロスを小さくしたいので下記のようになる        
        _iou_mask[cell_ind, a, :] = cfg.object_scale * (1 - iou_pred_cell_anchor)
        # IoUが最も高いアンカーボックスを選択
        _ious[cell_ind, a, :] = ious_reshaped[cell_ind, a, i]

        # マスク処理によってスケーリングを行い、ロス計算時の重要度を設定
        _box_mask[cell_ind, a, :] = cfg.coord_scale
        # 真のボックスをアンカーサイズに変換
        target_boxes[i, 2:4] /= anchors[a]
        # IoUが最も高いアンカーボックスを選択
        _boxes[cell_ind, a, :] = target_boxes[i]

        # マスク処理によってスケーリングを行い、ロス計算時の重要度を設定
        _class_mask[cell_ind, a, :] = cfg.class_scale
        # IoUが最も高いアンカーボックスを選択
        _classes[cell_ind, a, gt_classes[i]] = 1.

ロスの計算、先ほど導出したバウンディングボックス、IoU、クラスをスケーリングしてロスを算出しています。

            self.bbox_loss = nn.MSELoss(size_average=False)(bbox_pred * box_mask, _boxes * box_mask) / num_boxes
            self.iou_loss = nn.MSELoss(size_average=False)(iou_pred * iou_mask, _ious * iou_mask) / num_boxes

            class_mask = class_mask.expand_as(prob_pred)
            self.cls_loss = nn.MSELoss(size_average=False)(prob_pred * class_mask, _classes * class_mask) / num_boxes

Backward処理

下記でBackward処理が行われます。フレームワークを使う利点でforwardの処理さえ記述すればBackwardが自動な点が素晴らしいです。

    loss = net.loss
    loss.backward()

Optimize処理

Optimizerを下記で設定しています。今回はSGDを使用していますが他のOptimizerで比較しても良いと思います。

lr = cfg.init_learning_rate
optimizer = torch.optim.SGD(net.parameters(), lr=lr, momentum=cfg.momentum, weight_decay=cfg.weight_decay)

下記で最適化が終わります。

    optimizer.zero_grad()
    optimizer.step()

Save Model & Log処理

Save Model & Log処理は難しい部分はないので気になった部分を記述します。

下記でTensorboardへの書き込みをしています。ロス、学習率以外の値も知りたい場合は追記すれば良いですがモデルの出力などはまだ対応してないみたいです。

        if use_tensorboard and step % cfg.log_interval == 0:
            exp.add_scalar_value('loss_train', train_loss, step=step)
            exp.add_scalar_value('loss_bbox', bbox_loss, step=step)
            exp.add_scalar_value('loss_iou', iou_loss, step=step)
            exp.add_scalar_value('loss_cls', cls_loss, step=step)
            exp.add_scalar_value('learning_rate', lr, step=step)

学習率の変更を行っています。

        if imdb.epoch in cfg.lr_decay_epochs:
            lr *= cfg.lr_decay
            optimizer = torch.optim.SGD(net.parameters(), lr=lr, momentum=cfg.momentum, weight_decay=cfg.weight_decay)

高速化のための工夫

一例として並列処理の利用によるData augmentation処理の高速化を紹介します。

処理が設定されている箇所は下記で

imdb = VOCDataset(cfg.imdb_train, cfg.DATA_DIR, cfg.train_batch_size,
yolo_utils.preprocess_train, processes=2, shuffle=True, dst_size=cfg.inp_size)

preprocess_train内で下記の処理によって画像のサイズなどの変換を行っています。

im, trans_param = imcv2_affine_trans(im)

実際に画像のサイズが変換されているのが下記です。

    h, w, c = im.shape
    scale = np.random.uniform() / 10. + 1.
    max_offx = (scale - 1.) * w
    max_offy = (scale - 1.) * h
    offx = int(np.random.uniform() * max_offx)
    offy = int(np.random.uniform() * max_offy)

これらの処理をミニバッチ処理ごとに行っています。

                self.gen = self.pool.imap(self._im_processor,
                                          ([self.image_names[i], self.get_annotation(i), self.dst_size] for i in indexes),
chunksize=self.batch_size)