SushiGAN2 〜この寿司は実在しません〜

14076 ワード
GoogleCloudPlatform Python GAN DeepLearning Python テキストリンク

はじめに

2年ほど前，私はGANというものに出会って感銘を受け，Qiitaで『SushiGAN 〜人工知能は寿司を握れるか？〜』というちょっと釣りっぽいタイトルの記事を書きました．中身はChainerのexampleにあったDCGANの実装を寿司の画像で学習させてみるというもので，下のようなモード崩壊気味の寿司画像を作って満足していました．

あれから2年あまりが経ち，画像生成の技術は驚くべき発展を遂げました1．StyleGAN，そしてStyleGAN2が登場し，すぐには見分けがつかないほどの品質で1024x1024の人の顔の画像が生成できるようになりました [1]．そのインパクトは凄まじく，"This Person Does Not Exist"を皮切りに，GANで生成したリアルな偽物のサンプルをランダムに表示するWebサイトがちょっとした流行になりました．いくつか例を挙げると，次のようなものがあります．

私はこのシリーズのファンで，いつか自分のバージョンを作ってみたいと思っていたのですが，この度，ついに"🍣This Sushi Does Not Exist🍣"という寿司画像を生成するWebサイトを公開することができました．ご覧の通り，ただ画像を表示するだけの単純な代物ですが，これを見て興味を持った方が自分の"This Stuff Does Not Exist"を簡単に作ることができるように，この記事ではその作り方を簡単に紹介したいと思います．内容としては次のトピックを扱います．

  • Lightweight GANによる少ないサンプルでの学習
  • Google App EngineとFlaskを使ったWebサイトの実装

Lightweight GANの学習

最近のGANは，とてもリアルな画像を生成することができる反面，膨大な数のトレーニング画像と法外に高価な計算リソースを必要とすることが多いです．しかし最近，「Lightweight GAN」と呼ばれる匿名の著者による論文（現在，ICLR 2021のブラインドレビュー中）が，この問題をエレガントな方法で解決する方法を提案しました[2]．彼らは，skip-layer excitationモジュールとオートエンコーダー型の識別器という2つのシンプルな手法を用いて，100個程度の訓練サンプルと1台のGPUで数時間学習するだけで1024×1024の高解像度画像を生成することを可能にしました．おなじみのFFHQの顔画像はもちろん，ポケモンの画像まで悪くないクオリティで生成することができます．

画像は[2]より引用．

そして嬉しいことに，Phil Wang氏がLightweight GANの美しく使いやすい実装をGitHubで公開してくれました．これを使えば，たった2つのコマンドでColaboratory(または任意のGPU)上でLightweight GANを訓練することができます!

$ pip install lightweight-gan
$ lightweight_gan \
    --data ./path/to/images \
    --name {name of run} \
    --batch-size 16 \
    --gradient-accumulate-every 4 \
    --num-train-steps 200000

（詳しい手順はREADMEを参照してください．）

READMEには書いていない，寿司画像の生成に関する注意点を紹介します．

  • sushiscraperで250枚の寿司画像を集めました．
  • すべての画像を256x256にリサイズし，背景を黒く塗りました．
  • データ拡張にcutoutを入れると下のようなアーティファクトが発生することがわかったので，訓練時に--aug-types [translation]を設定しました．

訓練を始めて8時間後，DCGANとは比べ物にならないほど綺麗な寿司画像が生成されるようになりました．

寿司がヌルヌルと動く動画も作りました！

Webサイトの実装

"This Stuff Does Not Exist"のように生成画像をWebサイト上でに見せるために，Google App Engine（GAE）とFlaskを使ってThis Sushi Does Not Existを作りました．アクセスの度に画像を生成するとコストがかかるので，あらかじめ生成しておいた画像をGoogle Cloud Storage(GCS)に保存しておき，その都度ランダムに画像を取得して提示させるようにしました．GAEとGCS間の通信については，こちらの公式チュートリアルが非常にわかりやすかったです．This Sushi Does Not Existもこのチュートリアルをベースにしています．

main.pyとtemplates/index.htmlを以下のように変更してプロジェクトをデプロイすると，画像が提示されているのが確認できると思います．

main.py
import base64
import logging
import io
import os
import random

from flask import Flask, render_template
from google.cloud import storage

app = Flask(__name__)
CLOUD_STORAGE_BUCKET = os.environ['CLOUD_STORAGE_BUCKET']

gcs = storage.Client(project='PROJECT_NAME')
app.bucket = gcs.get_bucket(CLOUD_STORAGE_BUCKET)

@app.route('/')
def main():
    n = random.choice(range(384))
    blob = app.bucket.get_blob(f"{n:04}.jpg")
    buf = io.BytesIO()
    blob.download_to_file(buf)
    buf.seek(0)
    b64str = base64.b64encode(buf.getvalue()).decode("utf-8")

    return render_template("index.html", b64str=b64str)

if __name__ == '__main__':
    # This is used when running locally. Gunicorn is used to run the
    # application on Google App Engine. See entrypoint in app.yaml.
    app.run(host='127.0.0.1', port=8080, debug=True)

templates/index.html
<html lang="en" >
    <head>
        <meta charset="utf-8">
        <title>This Stuff Does Not Exist</title>
        <meta name="viewport" content="width=device-width,initial-scale=1">
        <meta name="description" content="This Sushi Does Not Exist">
    </head>
    <body>
        <div style="background: #000000; display: table; height: 100%; text-align: center; width: 100%;">
            <div style="display: table-cell; vertical-align: middle;">
                <div style="display: inline-block;">
                    <img src="data:image/jpg;base64,{{b64str}}" alt="sushi"/>
                    <div id="description" style="color: #FFFFFF;">
                        Find more details at... 
                    </div>
                </div>
            </div>
        </div>
    </body>
</html>

おわりに

本記事では，Lightweight GANとGAEを使って"🍣This Sushi Does Not Exist🍣"を作った方法を解説しました．「自分の集めた画像でGANを学習させる」のは寿司画像に限らずよく見かけるテーマですが，Lightweight GANを使えば相当の品質向上が期待できます．作ったモデルをWebアプリとして公開するのも簡単にできました．みなさんも，自分だけの"This Stuff Does Not Exist"を作ってみてはいかがでしょうか．

なお，ここではGCPを使う前提で説明しましたが，もちろんAWSでも同様のことができるはずです．AWSを使ったWebアプリケーションの作り方については，『AWSによるクラウド入門』がとても親切でおすすめです．

（1年ぶりにQiitaで記事を書きました．最近はブログをHippocampus's Gardenに移行して，英語の練習を兼ねてニッチな技術記事を書いています．よろしければこちらも見ていっていただけると嬉しいです．）

参考文献

[1] Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, Timo Aila. "Analyzing and Improving the Image Quality of StyleGAN." CVPR. 2020.
[2] Anonymous Authors. "Towards Faster and Stabilized GAN Training for High-fidelity Few-shot Image Synthesis." 2020.
[3] Using Cloud Storage  |  App Engine flexible environment for Python docs

  1. 昨年，『深層生成モデルを巡る旅(3): GAN』というGANのまとめ記事を書いたので，興味のある方はこちらもご覧ください． ↩