3. 少量のデータセットでの迅速なプロトタイピング

TAO Toolkit が提供するデータ拡張の機能と、NGC の学習済みモデルを併用することで、わずか 100 枚の画像で迅速にプロトタイプを作成することができます。

3.1 データ収集の課題

AI アプリケーションの開発において特にコストと時間がかかるのが、モデルのトレーニング用のデータセットの収集とアノテーションです。データが収集されるのを待っていては、開発スピードは鈍化してしまいます。また、モデルの学習に万全を期すには、十分な数と種類の画像が必要であり、どちらも不足しているような少量のデータセットでのみトレーニングしたモデルだと、パフォーマンスが低下する可能性もあります。

状況によっては、モデルの学習に十分な画像を収集できないこともあります。たとえば、プリント基板の組み立てラインで欠陥を検出する場合、この検出タスクは、品質を確保し、組み立てプロセスのエラーを発見するうえでとても重要です。ただし、プリント基板の組み立てで欠陥が発生する確率は非常に低いため、学習によってモデルの精度を高めるのに十分な画像を収集するには、数か月、あるいは数年かかることも考えられます。このように、異常を検出しようとする場合に十分なデータを収集することは、多くの業界で課題となっています。

この課題を克服するには、少量のデータセットに対して「データ拡張」と呼ばれる手法を用い、既存の画像をベースに種類やサイズを増やします。自前のデータセットを拡張することで、モデルの一般化を可能にする変動性とランダム性が加わり、これまで見たことのないデータに対するモデルの精度が向上します。より多くのデータを収集するよりも、大幅なコスト削減、期間短縮が可能です。

3.2 データの拡張

データ拡張は、既存のデータセットを利用し、空間や色のドメインで変換を施すことで、元の画像と類似し、かつ十分な差異を持つ画像を新たに作成する手法です。これにより、変動性の追加とモデルの一般化を図ります。これまでに、最も効果的な拡張方法を探るために多くの研究が行われてきました1。一般的な変換方式として、平行移動、回転、色変更などがあります。

少量のデータセットでトレーニングを行うと、そのモデルは、問題解決に必要な特徴を学習するのではなく、データのパターンを記憶し始めます。データ拡張によってデータセットのサイズを大きくすれば、データの複雑さが増し、記憶ではなく一般化せざるを得なくなります。これにより、学習用データセットに対する過剰適合が抑えられ、これまでに見たことのない画像に対するパフォーマンスが向上します。データ拡張は、さまざまな照明条件、位置、方向のオブジェクトに対応する可能性がある場合に特に有効です。

データ拡張を用いる際は、オフラインまたはオンラインで行います。オフライン拡張はトレーニングの前に行い、適用された変換によってストレージに新しい画像が作成されます。これにより、モデルが学習する一意の画像の枚数を制御でき、一般に、より少ないエポックでモデルを収束させることができます。オンライン拡張では、トレーニングに使用する各画像に対して、ランダム化された変換を動的に施します。画像が追加で保存されないため、余分なディスク容量は必要ありません。 また、変換のたびに一意の画像が作成されるため、常に新しい画像でモデルをトレーニングできます。オンライン拡張を用いたトレーニングでは、常に新しい画像が提示されるため、収束するまでに必要なエポック数が増える場合があります。

3.3 TAO Toolkit でのデータ拡張

TAO Toolkit は、オンラインとオフラインの両方のデータ拡張をサポートしています。オフライン拡張では、スペック ファイルを設定し、コマンド ライン インターフェイスから画像を生成します。この設定により、空間、色、ぼかしの拡張加減をカスタマイズします。オンライン拡張もカスタマイズ可能で、空間や色について、モデルのトレーニング時に適用する拡張範囲を指定できます。

TAO Toolkit でデータ拡張を行う場合に推奨されるのが、先にオフライン拡張でデータセットのサイズを大きくすることです。その後で、オンライン拡張でデータセットの複雑さがさらに増すようトレーニングを構成します。両方のデータ拡張を組み合わせることで、モデルに多種多様な画像を提示でき、後述の拡張タスクを実行した結果からもわかるように、モデルのパフォーマンスが向上します。

図 5. プリント基板の元画像
図 5. プリント基板の元画像
図 6. プリント基板の拡張後画像
図 6. プリント基板の拡張後画像
図 6. プリント基板の拡張後画像

TAO Toolkit でサポートされている拡張方式を次の表に示します。

空間 オフライン オンライン オフライン オンライン
回転 色相回転
反転 彩度変更
平行移動 コントラスト
傾き 輝度
ズーム 色変更
ぼかし
表 1.TAO における拡張

3.4 結論

データ拡張のメリットを示すために、欠陥のあるプリント基板の画像が 100 枚しかないデータセットを例にとります。このモデルのタスクは、プリント基板の画像から 6 種類の欠陥を検出することです。重要なパフォーマンス指標を mAP とし、モデルによって欠陥の周囲に配置されるバウンディング ボックスの精度を、グラウンド トゥルースと比較して評価します。

まず、データ拡張なしの 100 枚の画像のみでモデルをトレーニングしました。次に、オフライン拡張によって元の 10 倍と 20 倍のサイズのデータセットを作成し、比較のために、これらのデータセットでオンライン拡張ありのトレーニングと、オンライン拡張なしのトレーニングを行いました。

データ拡張なしの 100 枚の画像のみでトレーニングしたところ、テスト セットでの mAP は 36% にしか到達しませんでした。しかし、両方のデータ拡張を行った結果、モデルのパフォーマンスが大幅に向上し、約 79% の mAP を達成しました。このタスクは、画像が 100 枚という少量のデータセットの場合でも、データ拡張がモデルの品質向上に非常に有効であることを示しています。

図 7. 100 枚の画像での mAP
図 7. 100 枚の画像での mAP

ここまでのプロトタイプのモデルでは、100 枚の画像しか使用していません。そこで、さらに高い精度を達成するためにデータを増やすことにし、先ほどと同じタスクを 500 枚のアノテーション済み画像で実行します。データを増やしたことで、データ拡張を一切行わなくても mAP は 82% でした。さらに両方のデータ拡張を行った結果、mAP は 95% を超えました。

図 8. 500 枚の画像での mAP
図 8. 500 枚の画像での mAP

TAO Toolkit で提供されているデータ拡張を、さまざまなサイズのデータセットに活用してみることをお勧めします。TAO Toolkit は、少量のデータセットで機械学習アプリケーションを迅速に試作しようとする場合や、高精度のモデルのパフォーマンスをさらに改善しようとする場合に最適なツールです。

このタスクは、プリント基板欠陥データセット を用いて実行しました2。結果を再現するすべてのコードと手順は、TAO タスク GitHub リポジトリで公開しています。


1 Shorten, C., Khoshgoftaar, T.M. A survey on Image Data Augmentation for Deep Learning. J Big Data 6, 60 (2019). https://journalofbigdata.springeropen.com/articles/10.1186/s40537-019-0197-0

2 R. Ding, L. Dai, G. Li and H. Liu, "TDD-net: a tiny defect detection network for printed circuit boards," in CAAI Transactions on Intelligence Technology, vol. 4, no. 2, pp. 110-116, 6 2019, doi: 10.1049/trit.2019.0019.