私はニューラル ネットワークに基づく1枚画像超解像プログラムである “waifu2x” が登場した当時、これを利用させて頂き、機械学習の実験を繰り返して遊んでおりました。
その後、敵対的生成ネットワークと学習済み画像分類ネットワークVGG出力による知覚的損失関数を取り入れた “SRGAN” が登場し、これも何度も学習させたり改造したりしました。
以来、現在に至るまで様々な種類の深層学習による超解像を試して来ました。
そうした中で、私が利用するフレームワークはTensorFlow+TensorLayerからPyTorchに変わりました。
写実的な品質に関しては、以前作成したResidual of ResidualなネットワークのGANによる超解像ネットワークが最も優れていました。
ですが、あれはGenerator(画像生成器)とDiscriminator(真贋弁別器)という2つのネットワークをバランス良く学習させる必要があり、その調整が難しい事と、写実的にテクスチャーを生成しているものの高周波成分は本物ではなく架空のものであるので、動画生成ではフレーム毎に高周波成分が変化してしまうという短所がありました。
Relativistic GANという仕組みはGeneratorとDiscriminatorのバランスが崩れる問題を緩和してくれましたが、GANはネットワークを2つも学習させる必要がある事もあり、その後は違う方法を模索しました。
GANのDiscriminatorの代わりにAuto-Encoderのencoder出力を損失関数に組み込んだものも作成しました。
これは画質も良く、敵対的学習ではないので2つのネットワークのバランスを取る必要が無く、中々優れていましたが、チェッカーボード アーティファクトが発生する事があるという欠点がありました。
その後は損失関数をMSEやMAE、Sobelフィルターやその他の関数に置き換えて実験しました。
そしてこの度 “SSIM” (Structural SIMilarity / 構造的類似性)に基づく損失関数でネットワークを学習させ、ある程度の品質を達成致しました。
“GitHub” 上の “SRNet-DR” のページ:
https://github.com/ImpactCrater/SRNet-DR
SRNet-DRでは画像のノイズが綺麗に除去されています。
但し、これはGANではなく、更にノイズ除去を学習させている為、ランダムなパターンである毛髪などの細部は復元されていない事が分かります。
元の画像はiStockより使用ライセンスを購入致しました。
もちろん、これらのテスト画像は未学習画像です。
[SRNet-DRの特徴]
以前の記事と重複しますが、今回も概要を記載致します。
[1: Residual of Residual Network]
勾配消失に強い、残差学習を行います。
ショートカット接続のあるResidual of Residualの3重構造です。
上から下まで全て同一のResidual Block構成されており、shortcutに対して全てのConvolution層の出力との要素和を取っております。
Convolutionのカーネルは3×3で320chです。
Residual BlockはConv2d->Swish->Conv2d->Swishで、これを128段積み上げています。
[2: PixelShuffle]
画像の拡大はチェッカーボード アーティファクトが出難いように、最終段でPixelShuffle()x2です。
[3: Loss Function]
Generatorが生成した画像と正解の高解像度画像との “SSIM” (Structural SIMilarity / 構造的類似性)を計算し、その値をTensor要素毎に log(1+(a(1-SSIM))^2) として、これを損失値として学習させました。
SSIMは画像の局所的な比較で、平均(;色)と分散(;コントラスト)と共分散(;構造)を求めます。
MSE (Mean Squared Error / 誤差平方平均)だと学習結果はサンプル画像群の平均値寄りになり、MAE (Mean Absolute Error / 誤差絶対平均)は中央値寄りになり、0-1型損失関数だと最頻値(mode)寄りになります。
これらの中では最頻値寄りになった方が出力画像は鮮明になりますが、0-1型では微分が出来ないので学習出来ません。
そこで正規分布になるガウス関数を上下逆さにしたような関数を利用する事にしました。
log(1+(a(1-x))^2) がそれです。
SSIMだけでも画素毎の単純なMAEよりも鮮明な結果になるのですが、更にL0 lossに近付ける事でより鮮明になりました。
[4: Data Augmentation & Noise Reduction]
学習時のデータ拡張として、高解像度画像をランダムな位置でランダムな大きさにクロップします。
続いて所定の大きさまで縮小し、RGB画像をHSV形式に変換し、色相環をランダムに少しだけ回転させます。
ここまでがデータ拡張です。
その後、ランダムな強さでGaussian noiseを加えます。
そしてランダムな強さで圧縮を施して学習させます。
敢えてノイズを加え、圧縮して画質を落とした後、これを正解の高解像度画像に近付けるように学習させる事で、撮影時のノイズやJPEG画像や動画像のように圧縮ノイズが含まれた画像を拡大する際、ノイズを綺麗に除去出来るようになります。
JPEGでは固定の大きさのブロック単位で圧縮しますが、これがノイズ除去の柔軟性を失わせてしまうと考え、ブロックが目立ち難いWebP画像形式で圧縮を施しました。
WebPは不可逆圧縮でJPEGより軽量、高品質であり、可逆圧縮でもPNGよりもファイル サイズを小さく出来る優れた圧縮画像形式です。
結果は上出来で、RGBノイズが乗った画像や、強目の圧縮ノイズも綺麗に除去出来ます。
但しトレードオフとして、テクスチャーの高周波成分がやや失われてしまいます。
[5: Swish活性化関数]
活性化関数には性質の優れたSwish関数を用いております。
Swish(x)=xSigmoid(ax)
Sigmoid(x)=1/(1+e^(-a*x))
[6: Group Normalization]
ミニバッチ サイズが少ない場合に問題が起こり易いBatch Normalizationではなく、channleをグループ化して正規化するGroup Normalization層を用いています。
[7: 画像の読み込み]
学習させる画像群について、指定したディレクトリー以下にあるサブディレクトリーまで再帰的に検索して読み込みます。
画像はepoch毎にランダムにシャッフルされ、データ拡張されます。
全画像ファイル数がミニバッチの数で割り切れない場合には、不足分は既存の学習用画像の中からランダムに再度読み込まれます。
また、学習中にも画像を次々に追加出来ます。
指定したディレクトリー以下に、画像ファイルもしくは画像ファイルが入ったフォルダーを放り込むだけで、次のepoch開始時に読み込まれます。
[8: パラメーター数]
torchsummaryによる実測で、
Total params = Trainable params: 252,579,203
となり、約2億5千万パラメーターです。
[プログラムの実行について]
このプログラムで学習を実行するには18GB以上のメモリーが必要です。
CPUまたはNVIDIAのGPUによるCUDAで実行出来ます。
NVIDIA A100 GPU上でのメモリー消費量は、nvidia-smiによると、17,781MiB / 40,536MiBでした。
因みにGPUの消費電力は313Wとの事。
尚、低性能なCPUで実行する場合、充分な学習には数年以上の時間が掛かります。
今回私は、GCP (Google Cloud Platform)のCE (Compute Engine)でNVIDIA A100 GPUのVM Instance (a2-highgpu-1g)を利用させて頂きました。
しかも初期利用者限定の3万円分のクレジット付与により、実質無料で約3日間もA100 GPUをフル稼働させる事が出来ました。
GCPは非常にスケーラブルなプラットフォームであり、出来る事が個人のちょっとした事から大企業グループの大型プロジェクトまで対応しているので、その分初期設定はAWSのEC2よりも手間取ってしまいましたが、慣れれば何という事はありませんでした。
お陰で大きなモデルで22万枚強のサンプル画像を充分に学習させる事が出来ました。
ありがたい事です。
学習はミニバッチ数1枚で行い、1 step当たり0.34から0.35sでした。
以前のAWSのEC2のNVIDIA V100 GPU Instance(p3.2xlarge)の時はネットワークのチャンネル数が288chだったのを、今回のGCPのCEのNVIDIA A100 GPU VM Instance (a2-highgpu-1g)では320chに増やしたので計算量は約1.2倍ほどになっていますが、処理時間は1ステップ当たり0.63sから0.34sへと約0.54倍となりました。
つまりA100 GPUインスタンス(a2-highgpu-1g)はV100 GPUインスタンス(p3.2xlarge)よりも2.2倍強高速なようです。
尚、自宅のタワー型PCのIntel Core-i7 3770Tでは1 step当たり110s弱も掛かってしまいます。
ディープラーニングに於いて、NVIDIA A100 GPUはIntel Core-i7 3770T CPUよりも320倍も高速と言えそうです。
学習に使った画像は、自分で自宅及び屋外でカメラで撮影した沢山の静物(道具、工具、菓子、スウィーツ、立体パズル、植物、花、機械、電子回路基板、装飾品、素材、材料、部品、鉄塔、アンテナ塔、屋外階段など)、素材動画販売サイトから購入した多数の有料素材動画、無料素材動画配布サイト “Pexels.com” からダウンロードさせて頂いた多数の素材動画などです。
画像は膨大な枚数を全てチェックした中から、普遍的かつ問題の無いであろう22万枚強を選別して使用しています。
因みに2021年現在は、 “VQ-VAE” (Vector Quantised-Variational AutoEncoder)及び “VQ-VAE-2” や “Image Super-Resolution via Iterative Refinement” のように、非GANでありながら高精細な画像生成及び超解像が可能な手法が登場しております。
“VQ-VAE” はVAEの発展形です。
因みにVAEはAEと構造が似ていますが、出自が異なります。
VAEは潜在変数の確率分布として仮定された連続な多変量正規分布のパラメーターであるmu (平均)とsigma (分散)を学習させるものです。
ランダムなノイズを与えればランダムな画像を生成出来ます。
しかし現実世界では例えば犬と馬は連続的に変化するものではありませんので中間の曖昧な部分は不要です。
そこでVQ-VAEは潜在変数の確率分布をカテゴリカル分布として量子化して離散的な埋め込み表現にマッピングするようにしました。
中間が無いので生成画像は精細です。
画像生成時にはAutoregressiveモデルであるPixelCNNを利用して学習された潜在変数を生成し、それを元にdecoderで画像を生成します。
“Image Super-Resolution via Iterative Refinement” にある “Super-Resolution via Repeated Refinement” (SR3)では拡散確率モデルにより、正解画像に段階的にGaussian noiseが付加されて行き、最後には純粋なノイズ画像になる過程の逆過程、つまり徐々にノイズが減って鮮明化して行く過程を、入力された低解像度画像を条件にして繰り返し処理により学習させます。
ところで、もしFXトレードで儲かったら、上記の “SR3” のような拡散確率モデルの学習を試してみたいです。
しかしSR3は画像1枚を処理するのに数十回から100回の繰り返し処理が必要なのでかなり巨大な計算資源が必要になりそうです。
コメント