😇

超ニッチな二次元画像用リアルタイム超解像モデルを学習させた

2021/12/11に公開約7,100字

この記事は、創作+機械学習Advent Celendar 2021の11日目の記事です。

AIキャラクターの姿を表示する機能を作るにあたって軽量な超解像モデルが必要でした。既存の既存の軽量超解像モデルを調査したのですが、要件に合った速度の超解像モデルがなく結局自分で学習させた話です。
これにより、Talking Head Anime 2と併用して、1枚のイラストから、”高解像度"の表情を変えた画像を作れるようになりました。

この記事を読めばあなたの要件に合った軽量超解像モデルが見つかるかも知れませんし、超解像モデル学習時に役に立つかも知れません。Talking Head Anime 2を使ってAIのフロントエンドを作りたい人にも参考になるかもしれません。

作ったもの

AIキャラクターを表示するためのこういう画面を作りました。画面下のボックスにキャラクターに話してほしいことを入力しています。

こういう画面を作る場合、Unityを使ってVRMを扱うという手が王道かと思います。ただ、私はC#を書きたくない、Unityつらいなどと言った非常に正当な理由により、他の方法を探すこととしました。

Talking Head Anime 2を使えば、次の図の様に1枚の正面画像の表情を変えることができるため、ある程度AIキャラクターの表示に使えます。ただ、256x256の画像しか出力できないため、AIキャラクターの表示には小さすぎますそこで、超解像でなんとかする方向で考えることとしました。

超解像は、一言でいうと画像をきれいに拡大する技術です。
私の用途では、リアルタイムで表情を生成して、リアルタイムで拡大したかったので、かなり軽量で高速に動く超解像モデルが必要でした。
また、AIキャラクターは二次元キャラクターであるため、二次元画像用がきれいに拡大できる超解像モデルが必要でした。

軽量超解像の選択肢

まずは、既存のモデルで使えるものがないか調べてみました。
既存の軽量解像度の選択肢としては、次のようなモデルがあります。()内に書かれている速度は
256x256の画像を4倍に拡大した場合の速度で、RTX3070で測定しました。

  • Fast SRGAN(30fps)
  • Realtime Super Resolution(10fps)
  • Real ESRGAN for Anime(10fps)

Fast SRGAN

このモデルは、30fps近い速度で拡大ができ、リアルタイムで活用できる超解像モデルだと思います。実写の場合は高橋かずひとさんの検証によると、結構きれいに拡大できています。

アニメ画像の場合はかなりぼやけており、どちらかというと実写画像で役立つモデルかと思います。

私の場合、二次元画像をきれいに拡大したかったので、このモデルは合いませんでした。

Realtime Super Resolution

このモデルは10fpsと遅めな速度の上に、二次元画像の拡大品質はさほど高くありませんでした。
高橋かずひと先生の実験によると、実写画像を拡大するとこのようになるようです。

二次元画像に使うとこのような感じになります。
二次元画像拡大用途であれば結構つらそうです。

Real ESRGAN for Anime

Real ESRGANという非常にきれいに画像を拡大してくれるモデルがあります。
それのアニメ特化版はかなり軽量化されており、GPU上でなんと10FPSで動作します。
ただ、私の用途では最低でも30fpsは欲しかったので速度が少し足りませんでした。
拡大するとこのような感じになり、速度にシビアでない場合は、最高の選択肢になると思います。

軽量超解像モデルの学習

以上が既存の軽量超解像モデルでしたが、どのモデルも速度と品質のバランスがいい感じではありませんでした。そこで、自前モデルを学習することとしました。

まず、超解像の学習方法について説明します。
超解像では、2K画像くらいの比較的高解像度な画像が使われることが多いです。下記の画像のように高解像度画像から、32x32とか128x128という小さめのサイズ感のパッチを取り出して、それらのパッチを劣化させて学習データを作ります。そして、その学習データを用いて、学習したり評価します。
劣化処理と書いたのは一言に縮小と言っても、縮小の仕方にも種類があったり、ノイズが乗ったりする可能性もあるからです。
この劣化処理をどのように作成するかも非常に重要です。現実には複雑な劣化をするので、その複雑な劣化をうまく再現できなければノイズに弱いモデルになっていまいます
Real ESRGANでは、劣化処理を非常にうまく扱っています。

学習自体は、見栄え重視の超解像では、GANによる学習をすることが多いです。
GANの説明は省きますが、GANを用いた超解像では、Generatorに相当するものが超解像モデルで、Discriminatorは、超解像モデルが拡大画像か、本物の拡大画像化を当てるモデルとなります。これらのモデルが騙す、騙さないことを学習し合うことでGeneratorがきれいな拡大画像を生成できるようになります。
超解像をするGeneratorは、Discriminatorに本物の高解像度画像か、超解像された画像化を当てられない画像を作ることを学習し、Discriminatorはそれらを区別することを学習するイメージです。

色々と試行錯誤

軽量なモデルをベースに増やす

まずは、軽量なモデルをベースに改善する方針ではじめようと、先程紹介したFast SRGANをベースに改良する方向で始めました。
Fast SRGANはSRGANがベースのネットワークです。
ディープラーニングのモデルでは類似した構造を1ブロックとして積み上げることで深さを稼ぐモデルが多く見られますが、SRGANも同様です。
SRGANはResidual BlockというResnetで使われているブロックを積み上げている構造のモデルです。下記は論文中に書かれているモデル構造です。

このResidual Blockを更に軽量なInverted Residual Blockに置き換えたものがFast SRGANです。Inverted Residual Blockとは、MobileNetV2というスマートフォンでそこそこ早く動かせるモデルで採用されているブロックで、速度向上を見込んで採用したのだと思います。
Fast SRGAKはDIV2Kという実写画像のデータセットで学習されているため、これを二次元画像データセットで学習すればそこそこいい精度が出るものだと思い、Danbooru Dataset 2020のサブセット3000枚ほどで学習させてみました。1日学習させてもFast SRGANをそのまま使うよりはましなモデルにしかなりませんでした。
二次元画像を拡大した結果はこのとおりです。
びみょーにきれいになっていますが、ほぼ変わりません。

性能が高いモデルをベースに削る

Fast SRGANをベースに改善していくと道のりが長そうだったため、重量なモデルをベースにモデルを削っていく方法に切り替えました。
具体的には、先ほど紹介したReal ESRGANのアニメ版をベースにモデル削っていくこととしました。
先程説明しましたが、超解像の学習では、Discriminatorや劣化画像の生成方法も重要でReal ESRGANでは、これらにも最新の手法が取り入れられていたため、そういう観点から見ても早くゴールにたどり着きそうでした。

また、データセットをDanbooru 2020データセットの一部ではなく、Danbooruから高解像度画像を3000枚ほど取ってきたものを使いました。(highresタグのものを収集しました。)

Real ESRGANは、ESRGANで提案されたモデルをベースとしています。ESRGANで提案されたモデルは、RRDBというブロックを積み上げたネットワークで、RRDBはResidual Blockより重い代わりに表現力が高いブロックです。Real ESRGANのアニメ版では、RRDBのブロックを6個つみあげていて、かくRRDBのブロックの特徴抽出能力に相当する次元数は64次元でした。
このブロック数と次元数を減らしていくつか実験してみたところ、3ブロック、32次元くらいまで下げると25fpsになり、画質は落ちるとはいえそこそこ見栄えがいいものとなりました。
下記がずんだもんを拡大した結果です。

このモデルはそこそこいい感じになので、Hugging Faceで公開しています。

更に モデルの特徴抽出能力を16次元まで下げたモデルも試しましたが、かなり画質が悪くなり、このあたりがパラメータをいじる限界だと感じました。

論文を探し良さそうなモデルを作る

これ以上モデルを小さくすることはできなさそうだったので、別の方法を探すために、Real ESRGANやベースとなった論文の手法説明やAblation Studyを読みました。
論文のAblation Studyのセクションでは、提案手法の各パーツの性能への影響が調査されているため、速度重視のモデルを作る場合は参考になると思ったからです
ESRGANの論文にこのような図があり、超解像のモデル構造以上にDiscriminatorの構造や学習データセットの多様性が重要であることを示唆していました。

ESRGANのAblation Studyを受けて、元々軽量なSRGANでも、Discriminatorやデータ量が多ければ、そこそこいけるのでは?と思い、Real ESRGANのフレームワークでSRGANのモデルを学習させてみたところ、なんと256x256の画像を拡大した場合に50FPSでできて、そこそこ画質がいいモデルができあがりました。パラメータは6ブロックで特徴抽出能力を32としました。Fast SRGANより速い理由は、Fast SRGANは、CuDNNと相性が悪い、つまりNVIDIA GPUと相性が悪い可能性があります。原因調査はできていません。

このモデルもHugging Faceでデモと一緒に公開しているため、使ってみていただけると幸いです。
Hugging Face Spaceで公開しているアプリはたまに動かなくなるため、動かない場合はこのサンプルプログラムをローカルで動かしてみてください。

学習させてみて

リアルタイム超解像モデルの学習には論文調査や手法考案にもっと手間かかるのでは、と考えていました。実際は、4本程度の論文をつまみ読みする程度の知識で、二次元画像用のリアルタイムで動く超解像モデルを作ることができました。

タスクによりますが、モデルを学習するといっても、データセットを用意するだけだったり、今回の超解像モデルのようにモデルと学習手法の組み合わせやパラメータを変えるだけで学習できてしまうようなタスクもあります。そのため、普段モデルを学習させない人でもディープラーニングの基礎知識があるのであれば、学習させてみるのは楽しいかもしれません。

最も重要なのは計算リソースで、私は計算にはGoogle ColaboratoryのPro+プラン(月額5000円くらい)とRTX 3070を用いました。
これだけだと結構待ち時間が増えるので、スムーズに試行錯誤をするのであれば、あとRTX 3090があればいいと思っています!

私は持ってません(T_T)
誰か私のモデル開発を加速するために寄付してください!!笑 m(_ _)m

最後に

創作+機械学習Advent Celendar 2021の記事を二本も書かせていただきありがとうございました。

クレジット

下記画像のキャラクターのクレジット

イラスト:花兎*様

https://tyc.rei-yumesaki.net/material/illust/

キャラクター:フリー素材キャラクターつくよみちゃんの公式サイト

https://tyc.rei-yumesaki.net

下記画像のキャラクターのクレジット

音声読み上げ
VOICEVOX:ずんだもん

https://voicevox.hiroshiba.jp/

立ち絵
東北シリーズ公式

https://zunko.jp/con_illust.html

Discussion

ログインするとコメントできます