IMDN：軽量な超解像モデル
軽量な超解像モデルIMDNをさらに軽量化したモデルがあります。
onnxモデルもあります。RTX 3070環境で実行すると、256x256の画像を4倍にする場合、onnxruntime、Pytorch上で、約50FPSで動作します。超解像モデルの中では、非常に高速なモデルと言えます。
このモデルについてすこし理解を深めようと思います。
モデル構造
まず、モデル構造をざっくり論文を読んで理解します。
上記リポジトリで公開されているモデルは、IMDNモデルというものを更に軽量化したものです。IMDNモデルの構造は上記の図のようになっています。
基本的な構造は、SRGAN,ESRGANなどと似ていて、基本ブロックを積み上げて、入力画像と同解像度の特徴量マップを生成したのちに、アップサンプリングをするという流れです。この基本的なブロックが、SRGANが用いているSRResnetでは、Residual Blockなのに対して、IMDNでは、IMDブロックというブロックを用いています。
下記は実際に配布されているモデルの構造をtorchinfoでdepth=4で出力した結果です。
====================================================================================================
Layer (type:depth-idx) Output Shape Param #
====================================================================================================
IMDN -- --
├─Sequential: 1-1 [1, 3, 1024, 1024] --
│ └─Conv2d: 2-1 [1, 64, 256, 256] 1,792
│ └─ShortcutBlock: 2-2 [1, 64, 256, 256] --
│ │ └─Sequential: 3-1 [1, 64, 256, 256] --
│ │ │ └─IMDBlock: 4-1 [1, 64, 256, 256] 103,440
│ │ │ └─IMDBlock: 4-2 [1, 64, 256, 256] 103,440
│ │ │ └─IMDBlock: 4-3 [1, 64, 256, 256] 103,440
│ │ │ └─IMDBlock: 4-4 [1, 64, 256, 256] 103,440
│ │ │ └─IMDBlock: 4-5 [1, 64, 256, 256] 103,440
│ │ │ └─IMDBlock: 4-6 [1, 64, 256, 256] 103,440
│ │ │ └─IMDBlock: 4-7 [1, 64, 256, 256] 103,440
│ │ │ └─IMDBlock: 4-8 [1, 64, 256, 256] 103,440
│ │ │ └─Conv2d: 4-9 [1, 64, 256, 256] 36,928
│ └─Conv2d: 2-3 [1, 48, 256, 256] 27,696
│ └─PixelShuffle: 2-4 [1, 3, 1024, 1024] --
====================================================================================================
Total params: 893,936
Trainable params: 893,936
Non-trainable params: 0
Total mult-adds (G): 58.58
====================================================================================================
Input size (MB): 0.79
Forward/backward pass size (MB): 1233.13
Params size (MB): 3.58
Estimated Total Size (MB): 1237.49
====================================================================================================
Skip ConnectionはShortuctBlockという名前で実装されています。論文とは異なり、IMDBの各出力はシーケンシャルに次のブロックに入力されるようになっています。元の論文では、IICという、各ブロックの出力を結合？するような仕組みがありました。（論文には明記されていません）また、Conv-1もなくなっているようです。これらの変更はおそらく、さらなる軽量化の結果だと予想しています。Ablation Studyの結果を見てみても、IMDBの導入では、費用対効果が大きかったのに比べ、他のコンポーネントの効果はうすそうです。
このIMDブロックは、下記の図のように各層で16次元ずつ、すこしずつ特徴量を出力するようになっています。また、残りの48次元は、次のConvolutional Layerで処理されるようになっていて、完全にシーケンシャルな構造と層の間の出力を直接結合する方法の間を取っているようです。これで、パラメータ数を減らしつつ、PSNRやSSIMは向上しているようですが、なぜ向上しているかはわかりません。
実行とパフォーマンス測定
画像のコピーを含めた実行速度を測定してみたところ、pytorch, onnxruntimeともに256x256の画像を拡大する場合、20msくらいで動作します。超解像モデルの中では比較的高速に動作するモデルと言えます。また、拡大後の画像に関しても、そこそこ画質が良いものとなっています。
画質に関しては、かずひとさんのデモが参考になります。
実写画像では良さそうですが、アニメ画像ではあまりうまくいきませんでした。
あと、文字の超解像もうまくいきませんでした。
何に使えそうか？
リアルタイム性を活かして、検出などの認識系タスクと組み合わせるのが良いかも知れません。例えば、オブジェクト検出などで範囲を絞って、その範囲を拡大して、別の認識系タスクを実行するなどという方法が考えられます。もっと具体的な例としては、人の行動認識をしたい場合、あるサイズより小さい場合は超解像をかけてから行動認識をするといったこともできるかもしれません。
また、カラー画像以外の画像に視野を広げると面白いかもしれません。
また、上記でうまくいかなかった文字の超解像やアニメ画像に関しても、FineTuneするといい結果が得られるかも知れません。
Author And Source
この問題について(IMDN：軽量な超解像モデル), 我々は、より多くの情報をここで見つけました https://zenn.dev/xiongjie/articles/ba47fca40ec536著者帰属：元の著者の情報は、元のURLに含まれています。著作権は原作者に属する。
Collection and Share based on the CC protocol