[論文レビュー] Differentiable Data Augmentation with Kornia

Jian Shi, Edgar Riba+3 more|arXiv (Cornell University)|Nov 19, 2020

Advanced Neural Network Applications参考文献 11被引用数 8

ひとこと要約

この論文では、空間的および色の変換をニューラルネットワークの計算グラフに統合し、エンドツーエンドのバックプロパゲーションを可能にする、PyTorchネイティブで微分可能（differentiable）なデータオーグメンテーションフレームワーク、Kornia DDAを紹介する。PyTorchのautogradを活用し、オーグメンテーションパラメータの最適化を可能にすることで、GPUアクセラレートされ、再現可能で合成可能なデータオーグメンテーションパイプラインを実現し、高リソース環境下での優れた性能を示すとともに、微分可能オーグメンテーションポリシー探索のような新しい手法の実現を可能にする。

ABSTRACT

In this paper we present a review of the Kornia differentiable data augmentation (DDA) module for both for spatial (2D) and volumetric (3D) tensors. This module leverages differentiable computer vision solutions from Kornia, with an aim of integrating data augmentation (DA) pipelines and strategies to existing PyTorch components (e.g. autograd for differentiability, optim for optimization). In addition, we provide a benchmark comparing different DA frameworks and a short review for a number of approaches that make use of Kornia DDA.

研究の動機と目的

• 深層学習における非微分可能データオーグメンテーションの制限を解消し、オーグメンテーション操作を通過する勾配の流れを可能にすること。
• ニューラルネットワークトレーニングにシームレスに統合できる、PyTorchネイティブでモジュラーかつ合成可能なAPIを提供し、微分可能データオーグメンテーションを実現すること。
• 標準的なPyTorchトレーニングワークフローを用いて、バックプロパゲーションによる勾配最適化が可能なオーグメンテーションハイパーパrameter（例：明るさ、回転）の最適化を可能にすること。
• Kornia DDAを、TorchVision や Albumentations といった既存のライブラリと比較し、さまざまなハードウェアおよびデータ条件下での性能をベンチマークすること。
• 微分可能オーグメンテーションの実用的応用を、局所特徴記述子の学習や自動オーグメンテーションポリシー探索といった高度なタスクにおいて示すこと。

提案手法

• Kornia DDAは、PyTorchのautogradエンジンを用いて計算グラフに直接統合される微分可能レイヤー（例：RandomAffine、ColorJitter）としてデータオーグメンテーションを実装する。
• フレームワークは、PyTorchのnn.Moduleおよびnn.Parameterを用い、明るさやコントラストなどの微分可能パラメータをバックプロパゲーション中に最適化可能にする。
• グリッドサンプリング（grid_sample）などの微分可能演算を用いることで、空間的および色空間の操作における勾配計算を可能にする。
• フレームワークは2Dおよび3Dテンソルをサポートし、ボリュメトリックデータオーグメンテーション（例：CenterCrop3D、Affine3D）のための専用レイヤーを備える。
• オーグメンテーションパイプラインはモジュラーで、シリアル化可能なコンポonentsとして構築され、torch.saveおよびtorch.loadを用いてトレーニング済みモデルと併せて保存・ロード可能である。
• システムはCPU/GPU/TPUでのデバイス内計算を可能とし、PyTorchのランダムシード制御により再現性を確保する。

実験結果

リサーチクエスチョン

• RQ1 微分可能データオーグメンテーションは、PyTorchのネイティブトレーニングパイプラインに効率的に統合可能であり、オーグメンテーションパラメータのエンドツーエンド最適化を可能にするか？
• RQ2 さまざまなハードウェアおよびデータサイズ下で、TorchVision や Albumentations といった非微分可能フレームワークと比較して、Kornia DDAの性能はどの程度か？
• RQ3 微分可能オーグメンテーションは、勾配ベースの最適化によって局所特徴記述子の学習をどの程度向上できるか？
• RQ4 Kornia DDAは、Faster AutoAugment や MADAO のような効果的な自動オーグメンテーションポリシー探索を可能にするか？
• RQ5 Kornia DDAの微分可能性は、モデルとオーグメンテーションポリシーの共同最適化を可能にするような、新たなトレーニングパラダイムの実現を可能にするか？

主な発見

• 4台のGPUを用い、より大きな画像サイズ（例：512×512）で実行した場合、Kornia DDAはTorchVision や Albumentations よりも優れた推論速度を達成し、16.87秒の実行時間（Albumentations：18.99秒、TorchVision：16.12秒）を記録した。
• フレームワークはGPUメモリオーバーヘッドをほとんど発生させず、高性能トレーニングパイプラインへのデプロイに適している。
• Kornia DDAはコードの複雑さを20倍も削減し、カスタムPyTorchコードで約600行を要する複雑な微分可能変換を30行で実装可能にした。
• トレーニング後の最適化されたColorJitterパラメータは、明るさ：[0.8048, 0.8363]、コントラスト：[0.7030, 0.7323]、サチュレーション：[0.5999, 0.5976] であり、勾配ベースのチューニングに成功したことを示している。
• Kornia DDAは、微分可能オーグメンテーションを組み込んだモデルのエンドツーエンドトレーニングを可能にし、Faster AutoAugment や MADAO といったフレームワークが勾配降下を用いて最適なオーグメンテーションポリシーを学習可能にした。
• フレームワークは再現可能でデバイスに依存せず、シリアル化可能なオーグメンテーションパイプラインを提供し、モデルのポータビリティと実験の一貫性を向上させた。