background-eyecatch
eyecatch
公開日:2022/08/19更新日:2022/08/19

NAFNetで画像からノイズを除去するレシピ

---(0)
0

メニュー

概要

本レシピでは、NAFNetの技術概要を解説し、画像のノイズ除去を行う方法や、モデルのトレーニング方法をご紹介します。

NAFNet

Nonlinear Activation Free Network(以下、NAFNet)は、2022年4月に論文発表された画像復元技術です。

非線形活性化関数を使用しない本手法は、計算コストを大幅に削減しながら、ブレ除去用データセットであるGoProのベンチマークや、ノイズ除去用データセットSIDDのベンチマークで最先端のパフォーマンス(SOTA)を達成しています。

以下は、それぞれNAFNetによるブレ除去、ノイズ除去結果を示しています。

ブレ除去

ノイズ除去

画像復元技術の従来手法は、ディープラーニングの発展に伴い、精度が大幅に向上した一方で、システムが非常に複雑になるという問題がありました。

NAFNetでは、初めにシステムの複雑さをブロック間の複雑さと、ブロック内の複雑さに分解し、それぞれの複雑さを低減するベースラインを構築しています。

次に構築したベースラインに残存する非線形化関数をパフォーマンスを損なうことなく、別の関数に置換、または削除することにより、さらに単純化しNAFNetを構築しています。

ブロック間の複雑さ

従来手法を比較し、複雑でなくパフォーマンスに優れるアーキテクチャを実験的に選択しています。
以下では、それぞれのアーキテクチャを図示しており、NAFNetのベースラインでは最も単純なUNetを選択しています。


出典: Simple Baselines for Image Restoration

  • (a) Multi-Stage Architecture
    • 前段のステージの結果を後段のステージで改良。様々なサイズの特徴マップを複雑に接続
  • (b) Multi-Scale Fusion Architecture
    • ステージ内で、サイズの異なる特徴量を融合
  • (c) UNet Architecture
    • マルチステージでも、マルチスケールでもないシンプルなUNet

ブロック内の複雑さ

次に、ブロック内の内部構造を単純化するために最も一般的なコンポーネントである畳み込み、ReLUの構成(下図、(b) PlainNet's Block)に対して、近年SOTAを達成した技術で適用された手法を採用することで、下図(c) Baseline's Blockを構築しています。


出典: Simple Baselines for Image Restoration

具体的には、PlainNetから以下が追加・変更されています。

  • 学習の安定化、精度向上のため、Layer Normの追加
  • 精度向上のためReLUをGELUに置換
  • 計算効率とグローバルな情報を特徴マップにもたらすChannel-Attention(以下、CA)の追加

ここで、Restormer(上図 (a))で採用されているTransformarは、一部の研究でSOTAを達成するために必ずしも必要ではないとされていることや、複雑さが増すことから本研究では採用されていません。

NAFNetでは、上記ベースラインからさらにパフォーマンスを向上させるため、追加の改良を加えています。


出典: Simple Baselines for Image Restoration

一つ目は、GELUのSimpleGateへの置換です。

SOTAを達成する従来手法のいくつかがGated Linear Units(以下、GLU)を使用していることから、NAFNetでもGeLuからGLUへの置き換えを検討しています。

GLUは以下の式で表されます。

Gate(X,f,g,σ)=f(X)σ(g(X))Gate(X, f, g, σ) = f(X) ⊙σ(g(X))

ここで、XXは特徴マップ、ff, ggは線形変換器、σσは非線形活性化関数を表しています。このGLUを適用した場合、パフォーマンスが向上する可能性がありますが、ブロック内の複雑さも増加することが懸念されます。

このためNAFNetではGLUの簡素化のためGeLUを再検討します。
GeLUは以下の式で表されます。

GeLU(x)=xΦ(x)GeLU(x) = xΦ(x)

ここで、ΦΦは、標準正規分布の累積分布関数を表します。

上記GLUとGeLUと比較し、GeLUは、GLUのff, ggを恒等関数に、σσΦΦとした特殊ケースであると考えられます。
つまり、GLUはGeLUの一般化(活性化関数の一般化)であると捉えることができ、非線形活性化関数を置き換えることができると推察しています。

さらにGLU自体にも非線形性が含まれておりσσを削除してもGate(X)=f(X)g(X)Gate(X) = f(X) ⊙g(X)に非線形性が含まれているとしています。

ここから、特徴マップをチャネル次元で2つに分割し、乗算したSimpleGateを提案しています。


出典: Simple Baselines for Image Restoration

Simple Gateの式は非常にシンプルな以下の式で表されます。

SimpleGate(X,Y)=XYSimpleGate(X,Y) = X⊙Y

このSimpleGateは計算効率を向上させながら十分な精度が確保できることが実験から確認されています。

二つ目はCAのSimplified Channel Attention(以下、SMA)への置換です。

一つ目の置換により、残る非線形活性化関数は、CAのSigmoidとReLUのみとなっています。
さらなる計算効率の向上のためこれらの置換を検討しています。

まず、CAでは空間情報を絞り込みmultilayer perceptualを適用しattentionを計算しています。このようなCAは以下の式で表されます。

CA(X)=Xσ(W2max(0,W1pool(X)))CA(X) = X * σ({W_2}max(0, {W_1}pool(X)))

ここで、XXは特徴マップ、poolpoolは、空間情報をチャネルに集約するグローバル平均プーリングを表しています。σσは非線形活性化関数を表し、Sigmoid,W1,W2Sigmoid, W_1, W_2ではReLUが採用されています。

ここで、CAを一つの関数としてみなし、ΨΨと置き換えると以下のように表されます。

CA(X)=XΨ(X)CA(X) = X * Ψ(X)

この時、GLUの一般化同様、グローバルな情報の集約、チャネル情報の相互作用という2つの重要な役割を失わない範囲で単純化された以下のSCAに置き換えます。

SimplifiedChannelAttention(X)=XWpool(X)Simplified Channel Attention(X) = X * W pool(X)

SimpleGate同様、単純化によるパフォーマンスの損失がないことが実験から確認されています。


出典: Simple Baselines for Image Restoration

上記の通り、NAFNetでは、徹底した単純化により非線形活性化関数を排除し従来手法より大幅に計算量を削減しながら精度を向上させSOTAを達成しています。

以下はGoProデータセットを用いたブレ除去の比較結果です。
PSNR, SSIMはそれぞれ画質評価指標を、MACsは計算量を示しています。

計算量が低いにもかかわらず、高い精度を示していることが見て取れます。

それでは、実際にNAFNetを用いてブレ除去などの実行、および、モデルのトレーニングを動かしていきます。

学べること

  • NAFNetの技術概要
  • NAFNetを用いたDenoise方法
  • NAFNetモデルのトレーニング方法

始める前のスキルセット

  • Pythonの基本的な知識を有する人
  • 機械学習の基本的な知識を有する人
  • 「画像のノイズ除去」に興味を持っている人

実務活用例

  • 撮影画像の手振れ除去による鮮明化
  • 撮影画像のノイズ除去による鮮明化

キッチン

開発環境

  • Google Colaboratery
  • 下記の手順に従いGPUモードを利用してください
    • 「ランタイム」→「ランタイムのタイプを変更」→「ハードウェアアクセラレータ」をGPUに変更

ツール

  • Python3
  • Numpy, Pytorch等のライブラリ

データセット

調理

テスト

ブレ除去用データセットとして何のデータセットを使用しましたか?

NAFNetが、ReLU、GELUの代わりに使用している活性化関数は何ですか?

本レシピで実行していないタスクは何ですか?

おすすめレシピ