DNN の 2 倍の高速化を解き放つ: 統合ニューラル ネットワークを使用してわずか 1 分でモデルを変換

インテグラル ニューラル ネットワーク (INN) は柔軟なアーキテクチャであり、一度トレーニングすると、微調整することなく任意のユーザー定義のサイズに変換できます。音波 (音楽) は任意のサンプリング レート (音質) でサンプリングできるため、INN はデータとパラメーターの形状 (DNN 品質) を動的に変更できます。

TheStage.aiのチームは、IEEE/CVF CVPR 2023 カンファレンスで論文「Integral Neural Networks」を発表しました。この論文は、イベントで取り上げられたわずか 12 件の「賞候補」論文のうちの 1 つとして認められました。 INN は、連続パラメータと積分演算子を組み合わせて基本層を表す新しいクラスのニューラル ネットワークです。推論段階で、INN は連続重みからの離散サンプリングによって標準的な DNN 表現に変換されます。このようなネットワークのパラメーターはフィルターおよびチャネルの次元に沿って連続的であるため、それらの次元に沿って再離散化するだけで微調整を行わずに構造化された枝刈りが可能になります。

この記事では、4 倍画像超解像度 EDSR モデルを INN に変換するプロセスを説明し、モデルの構造化された枝刈りを実現する方法を示します。最後に、INN を個別の DNN に変換し、効率的な推論のために Nvidia GPU にデプロイします。この記事は次のように進めていきます。

1. INN について簡単に紹介します。
2. 超解像タスク用のEDSRネットワークの概要。
3. 1 行のコードで積分 EDSR を取得する TorchIntegral フレームワーク アプリケーション。
4. INN の微調整を行わない INN の構造プルーニング (高速パイプライン)。
5. プルーニングされたモデルを Nvidia GPU にデプロイします。

詳細と最新情報については、次のリソースを確認してください。

INNプロジェクトサイト

INNプロジェクトのGithub

この記事のサポートコード

DNN のための INN を微調整せずにプルーニングする

INN の層は積分演算子に置き換えられますが、積分演算子の実際の評価には数値積分法を使用するための入力信号の離散化が必要です。 INN の層は、離散化後に古典的な DNN 層 (全結合、畳み込み) と一致するように設計されていることがわかります。

EDSR 4x 画像超解像度プルーニング

画像の超解像度は、既知または未知の劣化演算子を使用して画像を強化する必要があるよく知られたコンピューター ビジョン タスクです。劣化演算子としてバイキュービック ダウンサンプリングを使用する古典的な超解像度形式を検討します。

画像の超解像度タスクには、拡散モデルやトランスフォーマーに基づくハイエンド ニューラル ネットワークなど、多くのアーキテクチャを使用できます。このドキュメントでは、4x EDSR アーキテクチャに焦点を当てます。 EDSR アーキテクチャは、ResNet (多くの深層学習問題で広く使用されています) と最後に 4x アップサンプリング ブロックで構成されているため、私たちのデモに最適です。 EDSR の概略説明を次の図に示します。

EDSR アーキテクチャのプルーニングの詳細

フィルタまたはチャネル全体の削除を伴う構造化プルーニングは、EDSR の主要な構成要素として機能する残差ブロックに対して独特の影響を及ぼします。各状態は Conv -> ReLU -> Conv ブロックを入力に追加することで更新されるため、入力信号と出力信号には同じ数のパラメーターが必要です。これは、枝刈り依存関係グラフを作成することによって、TorchIntegral で効率的に管理できます。以下の図は、各残差ブロックの 2 番目の畳み込みが単一のグループを形成することを示しています。

残差ブロック内の 2 番目の畳み込みを枝刈りするには、各残差ブロック内のすべての 2 番目の畳み込みを枝刈りする必要があります。ただし、より柔軟な設定のために、すべての残差ブロックで最初の畳み込みのフィルター、つまり 2 番目の畳み込みのチャネルをプルーニングします。

EDSR モデルから INN EDSR への変換

事前トレーニングされた DNN からの変換には、さらにスムーズな補間を備えた特別なフィルター チャネル順列アルゴリズムを利用します。順列アルゴリズムは、DNN の重みを連続関数からサンプリングされたように見せながら、モデルの品質を維持します。

 import torch import torchintegral as inn from super_image import EdsrModel # creating 4x EDSR model model = EdsrModel.from_pretrained("eugenesiow/edsr", scale=4).cuda() # Transform model layers to integral. # continous_dims and discrete dims define which dimensions # of parameters tensors should be parametrized continuously # or stay fixed size as in discrete networks. # In our case we make all filter and channel dimensions # to be continuous excluding convolutions of the upsample block. model = inn.IntegralWrapper(init_from_discrete=True)( model, example_input, continuous_dims, discrete_dims ).cuda()

統合グリッドのチューニング: DNN の構造化されたトレーニング後のプルーニング

統合グリッド調整は、ユーザー定義の数値に対してフィルターをサンプリングする必要があるパラメーター テンソルを (SGD 最適化の下で) スムーズに選択する操作です。フィルター/チャネル削除方法とは異なり、INN は補間演算により、複数の離散フィルターを組み合わせることができるフィルターを生成します。

 # Set trainable gird for each integral layer # Each group should have the same grid # During the sum of continuous signals # We need to sample it using the same set of points for group in model.groups: new_size = 224 if 'operator' in group.operations else 128 group.reset_grid(inn.TrainableGrid1D(new_size)) # Prepare model for tuning of integration grid model.grid_tuning() # Start training train(model, train_data, test_data)

実験中に、500 枚の画像が 4000 枚の画像のフル トレイン セット Div2k と同じ結果をもたらすことがわかりました。

パフォーマンス

結果として得られる INN モデルは、簡単に離散モデルに変換して、任意の NVIDIA GPU にデプロイできます。入力解像度 64x64 の RTX A4000 で 1 秒あたりのフレーム数 (FPS) を提供します。圧縮モデルでは、ほぼ 2 倍の高速化が実現します。枝刈りされた INN モデルを離散モデルに変換するには、次のコード行を使用できます。

 model = model.transform_to_discrete() # then model can be compiled, for instance # compilation can add an additional 1.4x speedup for inference model = torch.compile(model, backend='cudagraphs') 

結論

この記事では、CVPR2023 の賞候補論文「Integral Neural Networks」の概要を紹介しました。これは 4x EDSR モデルのトレーニング後の枝刈りに適用され、1 行のコードと 1 分間の統合グリッドの微調整でほぼ 2 倍の高速化を達成しました。

今後の記事では、INN のより多くのアプリケーションを紹介し、効率的なモデル展開についてさらに詳しく説明する予定です。乞うご期待：

INNプロジェクトサイト

INNプロジェクトのGithub

この記事のサポートコード

ご清聴ありがとうございました！