PyTorch 新たなクラスの物体検出をSSDでやってみる

今回は、血液の顕微鏡画像から細胞を検出するSSDモデルを作ってみたいと思います。

こんにちは cedro です。

先回、SSDの学習済みモデルを使った物体検出を行ってみましたが、物体検出できるのはあらかじめ学習した20クラスだけです。

新たなクラスの物体検出をするには、どうしたら良いのでしょうか。大量のアノテーション付きデータを用意してゼロから学習するしかないのでしょうか。

そんなことはありません。SSD学習済みモデルのパラメータを初期値として使う(過去の学習経験を活かす)と、少量のデータセットを学習するだけで、新しいSSDモデルを構築できるのです。これをファインチューニングと言います。

ということで、今回は、血液の顕微鏡画像から細胞を検出するSSDモデルをファインチューニングで作ってみたいと思います。

データセットを準備します

今回、使用するデータセットは、BCCD Dataset という血液の顕微鏡写真で、白血球、赤血球、血小板の3つについてバウンディングボックスのアノテーションデータが付いたものです。

データセットの仕様が、PASCAL Visual Object Classes ですので、PyTorch のSSDモデルで簡単に読み込むことが出来ます。

BCCD Dataset の検出イメージ

物体検出をした時のイメージは、こんな感じ。wbc が白血球、rbc が赤血球、platelets が血小板です。それにしても、なんともマニアックなデータセットですよね。

BCCD Dataset はこういった画像とアノテーションデータのセットが、全部で364個(trainval:292個、test:72個)しかない非常に小さなデータセットです。

BCCD_Dataset

実際に使用するのは、赤枠で囲ったBCCDフォルダーの部分です。非常にコンパクトなデータセットなので、今回はGithubのコードの中に含めてあります。

セットアップ

コードはGoogle Colabで動かす形にしてGithubに上げてありますので、それに沿って説明して行きます。自分で動かしてみたい方は、この「リンク」をクリックし表示されたノートブックの先頭にある「Colab on Web」ボタンをクリックすると動かせます。

まず、Githubからコードをコピーし、SSD学習済みモデルのパラメータをダウンロードします。ついでに、Google Colab に接続されているGPUの種類も確認しておきましょう。

 

ファインチューニングを行う前に、ダウンロードした学習済みモデルのパラメータを使ってVOCクラスの物体検出(20クラス)をやってみましょう。SSDネットワークをtestモードで定義し、学習済みモデルのパラメータを読み込みます。

そして、物体検出関数を定義し、サンプル画像(./data/person.jpg)の物体検出を行います。

左から順に dog, person, horse が検出されていて、数字はどれだけ自信を持って検出したかの確信度(%)を表しています。

そして、BCCDデータを読み込んで表示させてみましょう。

ファインチューニング

SSD学習済みモデルのパラメータを初期値として、BCCDデータセットを使ってファインチューニングを行います。初期設定の項目は以下の様です。

  • dataset : 学習に使うデータセット
  • basenet : ゼロから学習する場合に使用するパラメータ(不使用)
  • batch_size : バッチサイズ
  • resume : ファインチューングする場合に使用するパラメータ
  • max_iter : 学習回数(iter)
  • num_workers : 並列処理数
  • cuda : GPUの使用有無
  • lr : 学習率
  • lr_steps : 学習率を段階的に下げる場合の iter ポイント(不使用)
  • momentum : 最適化手法(momentumSGD)のパラメータ
  • weight_decay : 最適化手法(momentumSGD)のパラメータ
  • gamma : 学習率を段階的に下げる場合の学習率の低減率(不使用)
  • save_folder : 学習済みパラメータの読み込み・書き込みディレクトリ

ゼロから学習する場合は resume を指定せず、その場合 SSD は basenet で指定したパラメータをロードします。今回はファインチューニングなので resume を指定し、その場合SSDは resume で指定したパラメータをロードします。

学習率 lr は最初大きくし学習進行に合わせて段階的に下げると学習が安定し易いので、そのためのコードがありますが、実際には使用していません。また、今回はファインチューニングなので、パラメータの初期化は行いません。

それでは、学習を500iter実行します。GPUがP100の場合で、約10分ほど掛かります。

学習が完了したら、学習済みパラメータを使って物体検出をしてみましょう。

SSDをtestモードで定義し、先程ファインチューニングしたパラメータ ./weights/BCCD.pth をロードします。そして、BCCDのテストデータの画像を関数 detect() で物体検出を行います。

学習データがたった292個、学習回数も僅か500iter ですが、ちゃんと物体検出出来ています。もう1つやってみましょう。

では、また。

2021.5.12 内容をGoogle Colab バージョンに更新

3 件のコメント

  • すみません! cedro様

    ご丁寧なご回答ありがとうございます。

    せっかくご回答いただきましたのに、今日の今日まで cedro様の回答にすら気づかず、、、

    githubにデータあげてくださってありがとうございます!

    なんとか四苦八苦しながら、clolabで動かせたので、これから学習したpthで画像認識をしてみます!

  • cedro へ返信する コメントをキャンセル

    メールアドレスが公開されることはありません。 * が付いている欄は必須項目です

    日本語が含まれない投稿は無視されますのでご注意ください。(スパム対策)

    ABOUTこの記事をかいた人

    アバター

    ディープラーニング・エンジニアを趣味でやってます。E資格ホルダー。 好きなものは、膨大な凡ショットから生まれる奇跡の1枚、右肩上がりのワクワク感、暑い国の新たな価値観、何もしない南の島、コード通りに動くチップ、完璧なハーモニー、仲間とのバンド演奏、数えきれない流れ星。