以下はとある会社のtechブログにする体で書いていたものですが、blog記事になる前に別の会社に転職してしまったので、忘備録代わりにこちらに記載しました。去年の頭頃の記事なので情報が古いです。
概要
研修でSTARGANを実装しました。アルゴリズムや学習の構成等を解説します。さらにDocker imageをダウンロードし、Dockerコンテナ内で学習させ、簡単なウェブアプリを使ってテストをしてみます。Docker imageはここからダウンロードできますので、ご興味があればお試しください。現状未定
動機
FaceAppというアプリがあります。
これは顔写真をいろいろなテーマ(メガネ、男、女、老化、若返り等)に変身させることができます。例えば、ジョージ・ブッシュを変身させたものが以下になります。
左上が元画像、右上が女性、左下がハイゼンベルグ(ハゲ＋メガネ＋髭)、右下が若返りです。ディープラーニングでよく見られるスタイル変換のような絵画的な変換ではなく、かなり自然に変換されていることがわかります。このアプリ色々なテーマに変身させることが目的ですが、こういった個人で楽しむ以外にも、この技術自体は要素技術として色々応用ができそうです。
例えば、
- 警察の犯罪捜査への応用(似顔絵、モンタージュの高性能化)
- 美容院(色々な髪型を自然な形で再現)
- バーチャルユーチューバー
etc.
そこで今回はこのFaceAppと似たようなことができる方法を探したところ、STARGANという方法がそれに近いのではと思い実装しました(STARGANでないとできないというわけではないので、実際は違うアルゴリズムかもしれません)。
どんなアルゴリズムか
STARGANはディープラーニングのアルゴリズムで、いわゆるドメインを学習します。ドメインとは、例えばアウディ、シボレー、三菱ならドメインは車、インコや鷹、鶏ならドメインは鳥といった風に、共通の特徴を持ったものの集まりです。GANを使ったドメイン学習を行う生成モデルの中で、おそらくはじめに有名になったアルゴリズムはおそらくpix2pixです。これは下図のような様々なドメイン画像ペアを学習し、変換を行うことができます。
https://github.com/phillipi/pix2pix
ただし、この手法は学習の際に対応する画像対が必要になります。例えば上図の白黒→着色を学習させるためには白黒画像と全く同じものが写っているカラー画像が必要になります。白黒→着色データを集めるのはカラー画像さえ手に入ればそれを画像処理してやることで対応する白黒画像を手に入れることができるのですが、昼→夜の変換や、衛生画像→マップ画像の変換の場合、それらの画像対は必ずしも簡単に得られるものではないと思います。つまり、タスクによってはデータセット取得の敷居が高くなってしまうでしょう。
それを解決したのがCycleGANです。CycleGANは一つのドメインを学習するのにGeneratorとDiscriminatorをそれぞれ２つずつ、計4つ用意します。CycleGANはGeneratorのLossに画像対の差分(MSE, L1等)で誤差関数を定義する代わりに、一度Generatorに通してドメイン変換した画像に対し、さらにもう一つのGeneratorに通して元ドメインに通すことで再構成し、これと元画像の差分(cycle-consistency lossと呼ばれています)を定義します。これにより対応する画像対がなくてもドメインを学習することができます。このアルゴリズムはPix2Pix2と同じように様々なタスクに適用することができ、有名なものとしては馬をシマウマに変えることができます。データセットは単純に馬の画像とシマウマの画像を集めただけであり、Photoshopで馬に縞々を塗るなどはしていません。
https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix
上記で挙げたPix2PixとCycleGANは様々なドメインに対応できるのですが、仕様上2ドメイン間のドメイン変換しか行なえません。更に複数のドメインに変換するためには、その分(ドメインが $k$ 個あったら、$k(k−1)$個のGeneratorを学習しなければならない)のネットワークを用意しなくてはならず、少々非効率です。その問題を解決したのが、今回紹介するSTARGANというアルゴリズムです。
STARGAN
STARGANはHong Kong University of Science & TechnologyのYunjey Choiらが作ったアルゴリズムで、一つのGeneratorで複数のドメイン変換を行うことができます。下図の論文中記載のネットワーク図では星型をしていますが、実際のGeneratorは特殊なものではなく、Pix2PixやCycleGANで使われているResNetです(U-net等でも大丈夫だと思います)。その代わり、Generatorの入力と、Discriminatorの出力にそれぞれ細工をしています。
https://arxiv.org/abs/1711.09020
左が旧来のネットワークで複数ドメインの変換を実装した図、右がSTARGANです。Generatorだらけの従来手法よりスリムでかっこいいです(主観)。
ネットワークの入出力
例えば髪の色を変換したいとします。黒髪、ブロンド、ブラウンの髪の人の顔画像を集め、それらをGeneratorに与えます。ただそれだけではなく、それらの画像に写っている人が髪の色なのかも同時に教えてやります。具体的にはバッチサイズ分のカラー画像データと、ドメイン情報を0(無し) 、1(有り)の2値で表したラベルを、それぞれ画像サイズ分引き伸ばし、さらに画像とRGBのチャネル次元に連結して入力とします。
上記の場合、バッチサイズが3、{黒髪、ブロンド、茶髪}の3クラスのドメインラベルなので、(バッチサイズ3, ドメインラベル3個+チャネル数3個, 画像高さ128ピクセル, 画像幅128ピクセル) = (3, 6, 128, 128) のshapeのtensorを作り、Generatorに入力することになります。
次にDiscriminatorですが、入力は通常通り画像がバッチサイズ分結合されたtensorを入力としています。出力は、下図のように2つの出力が出てきます。
一つはデータが本物か偽物かの確率で、通常のDiscriminatorと同じ出力です。もう一つは画像がどのドメインに含まれるかの確率で、つまり一般のクラス分類の「ようなもの」です。なぜ「ようなもの」なのかというと、終端はソフトマックス関数ではなくsigmoid関数なため、複数のクラスで1を取ることができます。例えばメガネドメインと金髪ドメインを学習させたDiscriminatorにメガネをかけた金髪の人が入力してされると、2つのクラスが1になります。
学習の構成
STARGANのネットワークの入出力がわかったので、次は全体の構成を通して学習の仕方を説明します。データの流れとしては以下のようになります。
Discriminatorの訓練にはデータセットの画像(Real)と、Real画像からGeneratorで生成したFake画像を入力した時、それがRealだったら1を、Fakeだったら0を返すようにloss関数を設定します。これはAdversarial Lossと呼ばれており、以下の式の右辺第一項にあたります。
同時にドメインクラス分類の結果が出てきますが、その結果はロジスティック回帰の結果を返すようにします。
Generatorは少し複雑で、 まずOriginlal domain(Real画像のドメインクラス)をバッチ次元でランダムシャッフルし、Target domain(Fakeクラス)を作ります。Real画像とTarget domainをGeneratorに与えてFake画像を作ります。その次にFake画像をDiscriminatorに流し、Real/Fake判定をしますが、この時はFakeを与えたとき正解とします。これは一般のGANの学習で用いられており、GeneratorがDiscriminatorを騙しています。(1)式のAdversarial Lossの右辺第二項に対応します。 また、もう一つの出力であるドメインクラス分類誤差のLoss関数も定義します。
更に、先程Generatorで生成したFake画像を更にGeneratorに与えます。このときクラス情報としてOriginal domainを同時に与えます。これにより生成したRec画像(再構成画像)は、もとのReal画像と全く同じものになるべき画像です。Rec画像のReal画像の差分をとることによりLoss関数を定義します。これはCycleGANのCycle-consistency lossと等価です。
数式で表すと下記のとおりです。
論文中では$\lambda_{cls}=1, \lambda_{rec}=10$でした。上記２つのLoss関数をDiscriminatorとGeneratorで交互に最適化することで学習させていきます。
ペナルティ項の付与
学習の安定化とハイクオリティな画像の生成のために、Wasserstein GANで使われる勾配ペナルティ項を足しています。
ここで$\hat{x}$はReal画像とFake画像のランダムなアルファブレンドです。論文中では$\lambda_{gp}=10$でした。具体的な手法は異なりますが、RealとFakeを混ぜる正則化のアプローチはPix2PixやCycleGANでも使われていました。これらではSTARGANと違いDiscriminatorの入力にRealとFakeをchannel次元で結合して入力しています。STARGANで用いられていたアルファブレンドのほうがデータ量が少なくて良い気がします。
学習してみる
ネットワークについて理解が済んだので、次はデータセットをダウンロード学習をします。STARGANの論文で記載のあるCelebAのデータセットをダウンロードします。CelebAは20万枚のセレブの顔画像が含まれています。それぞれに対応する40個のドメイン(attributeと呼ばれている)が画像すべてにアノテーションされています。
| attribute | 意味 |
|---|---|
| 5_o_Clock_Shadow | (朝剃ったひげが伸びて)夕刻に目立ってくるひげ |
| Arched_Eyebrows | 三日月眉 |
| Attractive | 魅力的 |
| Bags_Under_Eyes | 目の下のたるみ |
| Bald | ハゲ |
| Bangs | 前髪 |
| Big_Lips | 大きな唇 |
| Big_Nose | 大きな鼻 |
| Black_Hair' | 黒髪 |
| Blond_Hair | 金髪 |
| Blurry' | ぼやけ |
| Brown_Hair | 茶髪 |
| Bushy_Eyebrows | ゲジゲジ眉毛 |
| Chubby | 丸ぽちゃ |
| Double_Chin | 二重あご |
| Eyeglasses | メガネ |
| Goatee | 山羊髭 |
| Gray_Hair | 白髪 |
| Heavy_Makeup | 厚化粧 |
| High_Cheekbones | 頬骨 |
| Male | 男 |
| Mouth_Slightly_Open | 口が僅かに開いている |
| Mustache | 髭 |
| Narrow_Eyes | 細目 |
| No_Beard | あごひげ無し |
| Oval_Face | 楕円形の顔 |
| Pale_Skin | 青白い肌 |
| Pointy_Nose' | 尖った鼻 |
| Receding_Hairline | 生え際後退 |
| Rosy_Cheeks | バラ色の頬 |
| Sideburns | もみあげ |
| Smiling | 笑い |
| Straight_Hair | 直毛 |
| Wavy_Hair | ウェーブ髪 |
| Wearing_Earrings | イヤリング |
| Wearing_Hat | 帽子 |
| Wearing_Lipstick | 口紅 |
| Wearing_Necklace | ネックレス |
| Wearing_Necktie | ネクタイ |
| Young | 若い |
このように、たくさんのドメインクラスが定義されています。すべてを学習すると入力のチャネル数が多すぎてGPUに乗り切らないので、10クラスで学習したいと思います。選んだクラスは(Bald, Black_Hair, Blond_Hair, Brown_Hair, Gray_Hair, Male, Straight_Hair, Wavy_Hair, Young)です。GPUは1080Tiを一本使用することにし、画像サイズを複数の条件{128x128, 256x256, 512x512}の３つの条件で試しました。画像が大きくなると、それだけGPUのメモリを消費してしまいます。そのためGPUのVRAMの容量に合わせ各条件ごとにバッチサイズを変更しました。以下が大まかな実験条件です。
| 画像サイズ | バッチサイズ |
|---|---|
| 128x128 | 20 |
| 256x256 | 8 |
| 512x512 | 4 |
学習の経過
ここからの作業は実際にコードを触って試してみましょう。こちら(どちら?)にdockerイメージが有りますので、ダウンロードしてご試用いただけます(dockerはこちらでインストールしてください)。また、CUDA9.1がインストールされている必要があります。
StarGAN.tarというファイルがダウンロードできると思います、そのディレクトリの直下で、下記の通り実行します
docker load < StarGAN.tar
これにより、お使いの環境にdockerイメージがインストールされます。下記のコマンドでStarGAN.tarがあるか確認してみましょう。
docker images
無事インストールされている場合は、下記のコマンドでコンテナを作成し、コンテナに入ります。
docker run -p 8097:8097 -p 8098:8098 --name stargan -it stargan /bin/bash
コンテナに入ったら、Visdomという学習状況を可視化するツールを以下のコマンドで起動します。
run_visdom
画像サイズが128の条件で学習するには
train_128
と打つとプログラムが走ります。
はじめにGeneratorとDiscriminatorのネットワーク構造が出力され、その下に学習状況(エポック数, 経過時間, 進捗率, 各Loss, 現在Batch数/総Batch数)が出力されます。
更に、学習状況を可視化してみましょう。先程起動したvisdomにhttp://localhost:8097でアクセスすることができます。
各Lossのグラフと、Real画像, Face画像, Rec画像, ドメインラベルそれぞれを1にしたときの生成画像(all)が表示されます。Real画像はBatch数分画像が表示されています。先程は画像サイズ128x128の条件で学習させましたが、その他256x256、512x512の条件での学習も試すことができます。256x256のときは
train_256
512x512のときは
train_512
で実行できます。画像サイズが異なると学習の経過や変換の度合いも変わるようです。どのように学習が進むかは、ご自身で確認してみてください。
テスト
最後に学習したモデルでテストをしてみましょう。
run_test
と打つとwebサーバーが起動します。http://localhost:8098にアクセスすると、以下のようなシンプルなページが表示されます。
ファイルを選択から変換したい画像(ただし、人の顔が写っているものに限る)を選択します。今回は髪の毛を金髪にしてみましょう。Blondにチェックを入れ、transformボタンを押します。
うまくいきました。その他のパラメータでうまくいくのかに関しては、ご自身で確かめてみてください。ちなみにうまく行ったものだけ載せています笑。
デフォルトでは256x256で学習させたネットワークを使っています。128x128や512x512を使う場合は下のラジオボタンにて設定してください。
所感
解像度が小さいほど、割りと大域的な変化をするように見えました。例えば髪のない人にWaivy Hairを適用すると、低解像度の場合は髪らしきものが生えてくるのですが、高解像度の場合は頭の色が変わるだけでした。局所的な変化(髪の毛の色)はどの解像度でも現れました。これはDiscriminatorが大域的な差異を見つけられていないのではないかということが考えられます。事例を調べてみるとCycleGANで犬を猫にできるかを試している方がいるようです。犬と猫は馬とシマウマと違って形に大きく違うため、通常のPatchGANではうまく学習ができないようです。そこで、Discriminatorを大域成分と局所成分を両方見るようにネットワークを変更したところ、大きな変化にも対応できたそうです。今回の解像度の違いも、こういったアプローチで対応できるかもしれません。私自身、犬猫変換をkaggleのデータセットを使って試したことがあります。結論としては128x128の条件では割りとうまくいき(荒いですが)、以下のような変化ができました。
|
|
|
ただし、それ以上の画像サイズではうまく行かず、上記のQiitaと同じように恒等変換や反転(ネガポジ変換のような変換)を学習してしまいました。やはり大域的な情報が重要のようです。