StyleGAN2で属性を指定して顔画像を生成する

memo.sugyan.com

の記事の続き(？)。 ある程度の学習データを収集して学習させたモデルが出来たので、それを使って実際に色々やってみる。

StyleGAN2-ADA

前回の記事でも書いたけど、厳選した16,000枚の画像を使って StyleGAN2-ADA を使って生成モデルを学習させてみた。

github.com

これは StyleGAN2 から進化したもので、より少ない枚数からでも安定して学習が成功するようになっていて、さらにparameter数など調整されて学習や推論もより早くなっている、とのこと。

それまでのStyleGANシリーズはTensorFlowで実装されていたが、最近はPyTorchに移行しつつある？のか、今はPyTorch版が積極的に開発進んでいるようだ。 そういう時代の流れなのか…。

github.com

今回は最初に触ったのがTensorFlow版だったのでまだPyTorch版は使っておらず、本記事はすべてTensorFlowだけを使っている。

学習

256x256程度のサイズのものであれば、Google Colabでも適度にsnapshotを残しながら学習しておき 切断されたら休ませてまた続きから学習始める、を繰り返せば数日〜数週間でそれなりに学習できる感じだった。 512x512やそれ以上のサイズだとちょっと厳しいかもしれない。

mapping出力と生成画像

StyleGAN2学習済みモデルを使ったmorphing、latent spaceの探求 - すぎゃーんメモ の記事で書いたが、StyleGAN の generator は、"mapping network" と "synthesis network" の2つの network によって作られている。 実際に画像を生成するのは synthesis network の方で、前段の mapping network は乱数入力を synthesis network への入力として適したものに変換するような役割になっている。

どちらの入力も潜在空間(latent space)としてみなせるが、 synthesis network への入力(= mapping network の出力)の方が次元数も多く、より表現力を持つものになっている。 この dlatents (disentangled latents) と呼ばれるものを線形に変化させることでスムーズなmorphingを表現できることを前述の記事で確かめた。

ということはこの dlatents の中に生成画像の属性を決定させるような要素があり、例えば顔画像生成のモデルの場合は顔の表情や向きなどを表すベクトルなどが存在しているかもしれない。 というのが今回のテーマ。

生成画像の属性推定結果から潜在空間の偏りを抽出

今回試したのは、以下のような手法。

• 生成モデルを使ってランダムに数千〜数万件の顔画像を生成
  • このとき、生成結果とともに dlatents の値もペアで保存しておく
• 生成結果の画像すべてに対し、顔画像の属性を推定する
• 推定結果の上位(または下位)数%を抽出し、それらを生成した dlatents たちの平均をとる

例えば表情に関する属性の場合、各生成結果の画像の「笑顔度」のようなものを機械的に推定し(もちろん手動で判別しても良いが、めちゃめちゃ高コストなので機械にやってもらいたい)、それが高scoreになっているものだけを集めて それらを生成した dlatents たちの平均値を計算する。 その値は、あらゆる顔画像を生成する dlatents の平均値と比較すると、笑顔を作る成分が強いものになっている、はず。

今回はまず学習済み生成モデルを使って適当な乱数から 20,000件の顔画像を生成し、それらに対して各属性を推定し、その結果で上位(または下位)0.5% の 100件だけを抽出するようにしてみた。 ものによってはもっとサンプル数が多く必要だったり、もっと少なくても問題ない場合もあるかもしれない。

表情推定

顔画像から表情を推定するモデルは幾つかあったが、今回は pypaz を利用した。

github.com

表情推定だけでなく、keypoint estimationやobject detectionなど様々な視覚機能を盛り込んでいて、それらを抽象化されたAPIで使えるようにしている便利ライブラリのようだ。

ここでは表情推定の部分だけを使用した。

import pathlib
from typing import Dict, List

import dlib
import numpy as np
import paz.processors as pr
from paz.abstract import Box2D
from paz.backend.image import load_image
from paz.pipelines import MiniXceptionFER


class EmotionDetector(pr.Processor):  # type: ignore
    def __init__(self) -> None:
        super(EmotionDetector, self).__init__()
        self.detector = dlib.get_frontal_face_detector()
        self.crop = pr.CropBoxes2D()
        self.classify = MiniXceptionFER()

    def call(self, image: np.ndarray) -> List[np.ndarray]:
        detections, scores, _ = self.detector.run(image, 1)
        boxes2D = []
        for detection, score in zip(detections, scores):
            boxes2D.append(
                Box2D(
                    [
                        detection.left(),
                        detection.top(),
                        detection.right(),
                        detection.bottom(),
                    ],
                    score,
                )
            )
        results = []
        for cropped_image in self.crop(image, boxes2D):
            results.append(self.classify(cropped_image)["scores"])
        return results


def predict(target_dir: pathlib.Path) -> Dict[str, np.ndarray]:
    results = {}
    detect = EmotionDetector()
    for i, img_file in enumerate(map(str, target_dir.glob("*.png"))):
        image = load_image(img_file)
        predictions = detect(image)
        if len(predictions) != 1:
            continue

        print(f"{i:05d} {img_file}", predictions[0][0].tolist())
        results[img_file] = predictions[0][0]

    return results

paz.processors として定義されたDetectorが、dlibで顔領域を検出した上でその領域に対し paz.pipelines.MiniXceptionFER によって表情を推定した結果を返してくれる。 MiniXceptionFER から返ってくるのは ['angry', 'disgust', 'fear', 'happy', 'sad', 'surprise', 'neutral'] の7 classesでの分類結果。

この結果で happy が1.0に近いものが得られたなら、それは確信度高く笑顔である、ということなので、その顔を生成した dlatents を集めて平均値を算出し、全体平均からの偏りをvectorとして抽出した。

これをmapping出力に加えていくことで、ランダムな生成顔画像も笑顔に変えていくことができた。

顔の向きや髪型も多少は影響を受けているが、概ね顔の特徴はそのままで主に口と目？あたりだけが変化している。 最初から笑顔だったものはより笑顔に、無表情だったものも口角が上がるくらいにはなっている。 口が開いて前歯が見えるようになったりするのも興味深い。

ちなみに他の表情に関しては、同様のことをやっても「怒っている顔」や「悲しい顔」のようなものは作れなかった。 まず、今の生成モデルの学習のために収集し厳選したデータは大部分が笑顔か無表情の顔画像であったため(悲しい顔ばかり載せているようなアイドルは居ない)、生成される画像も多くはそのどちらかであり、それ以外の表情の画像はほぼ生成されない。 ので、機械推定した結果もそれらの表情を強く検出するものはなく、笑顔のようにベクトルを抽出することは難しいようだ。

無表情化の場合。元が笑顔のものが真顔になる程度の変化は一応ありそう。 真ん中上段の子は前髪がだいぶアレだが…。

悲しみ。少し笑顔が消えて眉が困ってそうな感じになっているようには見える。 が、とても微妙…。

怒り。何故か顔の向きばかり変化してしまっているが、結局表情はほとんど変化が無いようだ。

・結果まとめ

笑顔に関しては概ね上手く抽出できた。 その他の表情についてはほとんど良い結果にならなかったが、学習データに使う顔画像の表情がもっと豊富にあればそういった顔画像も生成できて抽出が可能になると思われる。

顔姿勢推定

顔の向きも推定して同様のことをしてみる。 機械学習モデルでも顔角度の推定できるものありそうだが、今回は dlib で検出したlandmarkの座標から計算する、というものをやってみた。

詳しくは以下の記事を参照。

learnopencv.com

ほぼこの記事の通りに実装して、入力画像から yaw, pitch, roll のEuler anglesを算出する。 参照している3Dモデルの座標や数によって精度も変わってきそうだが、とりあえずは上記記事で使われている6点だけのものでそれなりに正しく角度が導き出せるようだった。

import math
import pathlib
from typing import List

import cv2
import dlib
import numpy as np


class HeadposeDetector:
    def __init__(self) -> None:
        predictor_path = "shape_predictor_68_face_landmarks.dat"

        self.model_points = np.array(
            [
                (0.0, 0.0, 0.0),  # Nose tip
                (0.0, -330.0, -65.0),  # Chin
                (-225.0, 170.0, -135.0),  # Left eye left corner
                (225.0, 170.0, -135.0),  # Right eye right corne
                (-150.0, -150.0, -125.0),  # Left Mouth corner
                (150.0, -150.0, -125.0),  # Right mouth corner
            ]
        )
        self.detector = dlib.get_frontal_face_detector()
        self.predictor = dlib.shape_predictor(predictor_path)

    def __call__(self, img_file: pathlib.Path) -> List[float]:
        image = cv2.imread(str(img_file))
        size = image.shape

        # 2D image points
        points = []
        rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        dets = self.detector(rgb, 1)
        if len(dets) != 1:
            return [np.nan, np.nan, np.nan]

        d = dets[0]
        shape = self.predictor(rgb, d)
        for i in [30, 8, 36, 45, 48, 54]:
            points.append([shape.part(i).x, shape.part(i).y])
        image_points = np.array(points, dtype=np.float64)

        # Camera internals
        focal_length = size[1]
        center = (size[1] / 2, size[0] / 2)
        camera_matrix = np.array(
            [[focal_length, 0, center[0]], [0, focal_length, center[1]], [0, 0, 1]],
            dtype=np.float64,
        )

        # Calculate rotation vector and translation vector
        dist_coeffs = np.zeros((4, 1))  # Assuming no lens distortion
        success, rotation_vector, translation_vector = cv2.solvePnP(
            self.model_points,
            image_points,
            camera_matrix,
            dist_coeffs,
            flags=cv2.SOLVEPNP_ITERATIVE,
        )

        # Calculate euler angles
        rotation_mat, _ = cv2.Rodrigues(rotation_vector)
        _, _, _, _, _, _, euler_angles = cv2.decomposeProjectionMatrix(
            cv2.hconcat([rotation_mat, translation_vector])
        )

        return [
            math.degrees(math.asin(math.sin(math.radians(a))))
            for a in euler_angles.flatten()
        ]

表情と同様に、値の大きなものを生成した dlatents の平均、値の小さなものを生成した dlatents の平均、を使って顔向きを変化させるベクトルを求める。

yaw (左右向き)。全体の大きさが変わってしまうので連続的に動くと不自然に感じてしまうが、ともかく顔の特徴はそのままに向きだけが変化しているのは観測できる。 顔の向きが変わっても視線は固定されている。

pitch (上下向き)。yawほど顕著に差は出ないが、それなりには変化する。 真ん中上段の子の前髪はやはりハゲやすいようだ…。

roll (傾き？)。これは首を傾けるように変化するはずのものだが、そもそも 学習データの前処理の段階 で正規化されているので傾いた画像が生成されるはずがない。 ので表情のときと同様に正しくベクトルを抽出できず、結果的に何故かyawと似たような動きになってしまっている。

・結果まとめ

yaw, pitch それぞれに関しては概ね上手く抽出できた。yawの方が学習データのバリエーションが多かったからか、より顕著に差が出るようだった。

髪領域推定 (顔解析)

次は髪色、髪の長さなどを変化させたい。 これらの属性を数値として得るには、まず髪の領域を抽出する必要がある。

ここでは、TensorFlowで動かせる学習済みモデルとして face_toolbox_keras を使用した。

github.com

これもpypazのように顔やlandmarkの検出など様々な機能があるが、その中の一つとして Face parsing がある。 元々は https://github.com/zllrunning/face-parsing.PyTorch で、それを移植したものらしい。 この Face parsing は入力の顔画像から「目」「鼻」「口」など約20のclassに各pixelを分類する。 髪の領域は値が 17 になっているので、その領域だけを抽出することで顔画像から髪の部分だけを取り出すことができる。

髪の領域だけ得ることができれば、あとは

• その面積で「髪のボリューム」
• 下端の位置・下端の幅で「髪の長さ」
• 画素の平均値で「髪色の明るさ」

などを数値化できる。 顔姿勢と同様に、値の大きなものを生成した dlatents の平均、値の小さなものを生成した dlatents の平均、を使ってベクトルを求める。

ボリューム。面積だけで計算しているのでちょっと頭の形が変になったりするかもしれない…。

長さ。ボリュームよりは自然な感じで長さが変化しているように見える。 真ん中上段の子の前髪はやはり(ry

明るさ。暗くするとみんな黒髪になるし、明るくすると茶髪や金髪など様々な明るい色になる。

・結果まとめ

髪の領域を推定することで、髪に関する属性を計算することができて長さや色などを変化させることができた。 もう少し頑張って上手く数値化できれば、前髪の具合や触覚の有無指定などもできるようになるかもしれない。

年齢 (上手くいかず)

これも他と同様、理論的には「顔画像から年齢を推定し、高い数値のものを生成した dlatents と 低い数値のものを生成した dlatents からベクトルを抽出」という感じで 童顔にしたり大人びた顔にしたりできると思っていたのだけど、そもそもの年齢の推定が全然正確にできなそうで断念した…。

• deepface
• Age and Gender Estimation (yu4u/age-gender-estimation)

など幾つかの学習済みモデルを使って年齢推定をかけてみたのだけど、どれも「東アジアの10〜20代女性」の学習データが乏しいのもあるのか(もしくは使い方が悪かった…？)、結果のブレが激しくて とても正しい年齢推定ができている感じがしなかった。

生成の学習に使ったデータからある程度は年齢ラベルつけたデータセットは作れるので、頑張れば年齢推定モデルを自前で学習させて より正確な推定ができるようになるかもしれないが… そこまでやる気にはならなかったので諦めた。

複合

とりあえずはここまでで

• 表情 (笑顔◎、無表情△)
• 顔角度 (左右◎、上下○)
• 髪 (長さ◎、明るさ◎)

といった属性については変化させるためのベクトルが抽出できた。 ので、複数を足し合わせたりすることもできる。

無表情 + 右上向き + 髪短く明るく

笑顔 + 下向き + 髪長く暗く

というわけで、記事の冒頭に貼った画像は元は同じ顔からこうして変化させて作ったものでした。

Repository

• https://github.com/sugyan/gan-playground

N番目の素数を求める

SNSなどで話題になっていたので調べてみたら勉強になったのでメモ。

環境

手元のMacBook Pro 13-inchの開発機で実験した。

• 2.8 GHz Intel Core i7
• 16 GB 2133 MHz LPDDR3

Pythonでの実装例

例1

最も単純に「2以上p未満のすべての数で割ってみて余りが0にならなかったら素数」とする、brute force 的なアプローチ。

import cProfile
import io
import pstats
import sys


def main(n: int) -> int:
    i = 0
    for p in range(2, 1000000):
        for q in range(2, p):
            if p % q == 0:
                break
        else:
            i += 1
        if i == n:
            return p
    raise ValueError


if __name__ == "__main__":
    n = int(sys.argv[1])
    with cProfile.Profile() as pr:
        for _ in range(10):
            result = main(n)
    print(result)

    s = io.StringIO()
    ps = pstats.Stats(pr, stream=s)
    ps.print_stats("main")
    print(s.getvalue())

Python3.9.1 で実行してみると

$ python3.9 1.py 1000
7919

...

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
       10    2.237    0.224    2.237    0.224 /Users/sugyan/dev/sugyan/nth-prime-benchmark/python/1.py:7(main)

1,000番目のものを求めるのにも 224 msほどかかっていて遅い。10,000番目すら求めるのが大変。

例2

求めた素数をlistに入れていき、それらで割れるかどうかだけ確認していく。いわゆる「試し割り法」(trial division)というらしい。

試し割り法 - Wikipedia

def is_prime(num: int, primes: List[int]) -> bool:
    for p in primes:
        if num % p == 0:
            return False

    primes.append(num)
    return True


def main(n: int) -> int:
    i = 0
    primes: List[int] = []
    for p in range(2, 1000000):
        if is_prime(p, primes):
            i += 1
        if i == n:
            return p
    raise ValueError

これだと10倍ほど速くなって1,000番目も 27 ms程度で出てくる。

$ python3.9 2.py 1000
7919

...

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
       10    0.015    0.002    0.275    0.027 /Users/sugyan/dev/sugyan/nth-prime-benchmark/python/2.py:17(main)

10,000番目だと 2,607 msくらい。

$ python3.9 2.py 10000
104729

...

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
       10    0.205    0.021   26.067    2.607 /Users/sugyan/dev/sugyan/nth-prime-benchmark/python/2.py:17(main)

例3

「エラトステネスのふるい」と呼ばれるものを疑似したもの。limit までの数に含まれる素数を篩にかけて列挙していって、N個以上あればN番目を返す、無ければlimitを倍にしていく。

def faked_eratosthenes(limit: int) -> List[int]:
    nums = [i for i in range(2, limit + 1)]
    primes = []
    while True:
        p = min(nums)
        if p > math.sqrt(limit):
            break
        primes.append(p)
        i = 0
        while i < len(nums):
            if nums[i] % p == 0:
                nums.pop(i)
                continue
            i += 1
    return primes + nums


def main(n: int) -> int:
    limit = 1000
    while True:
        primes = faked_eratosthenes(limit)
        if len(primes) > n:
            return primes[n - 1]
        limit *= 2

例2のものより多少速くなるが、それほど大きくは変わらない。

$ python3.9 3.py 10000
104729

...

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
       10    0.004    0.000   19.244    1.924 /Users/sugyan/dev/sugyan/nth-prime-benchmark/python/3.py:26(main)

エラトステネスの篩

下記記事が詳しいが、前述の例は「似非エラトステネスの篩」として言及されている。

zenn.dev

正しくは値をリストから削除するのではなくフラグで管理していく、とのこと。

def list_primes(limit: int) -> List[int]:
    primes = []
    is_prime = [True] * (limit + 1)
    is_prime[0] = False
    is_prime[1] = False

    for p in range(0, limit + 1):
        if not is_prime[p]:
            continue
        primes.append(p)
        for i in range(p * p, limit + 1, p):
            is_prime[i] = False

    return primes


def main(n: int) -> int:
    limit = 1000
    while True:
        primes = list_primes(limit)
        if len(primes) > n:
            return primes[n - 1]
        limit *= 2

こうすると(1回の list_primes の中では)リストのサイズ変更がなくなり領域の再確保やコピーもなくなり、倍数を篩によって除外するのも速くなる、ということ。

$ python3.9 eratosthenes.py 10000
104729

...

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
       10    0.004    0.000    0.440    0.044 /Users/sugyan/dev/sugyan/nth-prime-benchmark/python/eratosthenes.py:24(main)

$ python3.9 eratosthenes.py 100000
1299709

...

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
       10    0.092    0.009    7.978    0.798 /Users/sugyan/dev/sugyan/nth-prime-benchmark/python/eratosthenes.py:24(main)

10,000番目を求めるのも 40倍ほど速くなり、 100,000番目くらいでも 798 ms 程度で求められる。

Rustでの実装例

Pythonで満足したので次はRustで書いてみる。

試し割り法

pub fn trial_division(n: usize) -> u32 {
    let mut primes = Vec::with_capacity(n);
    primes.push(2u32);
    while primes.len() <= n {
        if let Some(prime) = (primes[primes.len() - 1] + 1..).find(|&m| {
            primes
                .iter()
                .take_while(|&e| e * e <= m)
                .all(|&e| m % e != 0)
        }) {
            primes.push(prime);
        }
    }
    primes[n - 1]
}

エラトステネスの篩

pub fn eratosthenes(n: usize) -> u32 {
    fn list_primes(limit: usize) -> Vec<u32> {
        let mut primes = Vec::new();
        let mut is_prime = vec![true; limit + 1];
        is_prime[0] = false;
        is_prime[1] = false;
        for p in 0..=limit {
            if !is_prime[p] {
                continue;
            }
            primes.push(p as u32);
            for i in (p * p..=limit).step_by(p) {
                is_prime[i] = false;
            }
        }
        primes
    }

    let mut limit = 1000;
    loop {
        let primes = list_primes(limit);
        if primes.len() > n {
            return primes[n - 1];
        }
        limit *= 2;
    }
}

アトキンの篩

全然知らなかったのだけど、エラトステネスの篩よりも速いアルゴリズムとして「アトキンの篩」というものがあるらしい。

Sieve of Atkin - Wikipedia

エラトステネスの篩の計算量 に対し こちらは になる、とのこと。

原理は正直全然わからないけど、

• 4 * x * x + y * y == n となる n の mod 60 が {1, 13, 17, 29, 37, 41, 49, 53} に含まれる
• 3 * x * x + y * y == n となる n の mod 60 が {7, 19, 31, 43} に含まれる
• x > y かつ 3 * x * x - y * y == n となる n の mod 60 が {11, 23, 47, 59} に含まれる

の3つの式において x と y の組み合わせの数が合計で 奇数個存在している 場合に、nを素数の候補とすることが出来て、それらから平方数を除いたものが素数として列挙できるようだ。 (ちょっと解釈が間違ってるかもしれない)

とりあえず効率はあまり考えずに見様見真似で実装してみた。

pub fn atkin(n: usize) -> u32 {
    fn list_primes(limit: usize) -> Vec<u32> {
        let mut primes = Vec::new();
        if limit > 2 {
            primes.push(2);
        }
        if limit > 3 {
            primes.push(3);
        }
        let mut sieve = vec![false; limit];
        for x in (1..).take_while(|&x| x * x < limit) {
            for y in (1..).take_while(|&y| y * y < limit) {
                {
                    let n = (4 * x * x) + (y * y);
                    if n <= limit && (n % 12 == 1 || n % 12 == 5) {
                        sieve[n] ^= true;
                    }
                }
                {
                    let n = (3 * x * x) + (y * y);
                    if n <= limit && n % 12 == 7 {
                        sieve[n] ^= true;
                    }
                }
                if x > y {
                    let n = (3 * x * x) - (y * y);
                    if n <= limit && n % 12 == 11 {
                        sieve[n] ^= true;
                    }
                }
            }
        }
        for r in (5..).take_while(|&r| r * r < limit) {
            if sieve[r] {
                for i in (1..).map(|i| i * r * r).take_while(|&i| i < limit) {
                    sieve[i] = false;
                }
            }
        }
        primes.extend(
            sieve
                .iter()
                .enumerate()
                .filter_map(|(i, &b)| if b { Some(i as u32) } else { None }),
        );
        primes
    }

    let mut limit = 1000;
    loop {
        let primes = list_primes(limit);
        if primes.len() >= n {
            return primes[n - 1];
        }
        limit *= 2;
    }
}

おまけ: GMP

多倍長整数の算術ライブラリとしてGMPがあり、これのRust bindingがある。

https://crates.io/crates/rust-gmp

Mpz::nextprime() というのを呼ぶと「次の素数」を求められるらしいので、N回実行すればN番目の素数が求められそうだ。

gmp::mpz::Mpz - Rust

pub fn gmp(n: usize) -> u32 {
    let mut mpz = gmp::mpz::Mpz::new();
    for _ in 0..n {
        mpz = mpz.nextprime();
    }
    mpz.to_string().parse().unwrap()
}

Benchmark

というわけでこのあたりでbenchmarkを取ってみると

$ rustup run nightly cargo bench

...

test bench_100000_atkin           ... bench:  13,437,614 ns/iter (+/- 3,614,150)
test bench_100000_eratosthenes    ... bench:  30,768,639 ns/iter (+/- 24,144,191)
test bench_10000_atkin            ... bench:     858,282 ns/iter (+/- 724,131)
test bench_10000_eratosthenes     ... bench:   1,783,792 ns/iter (+/- 269,701)
test bench_10000_gmp              ... bench:  19,331,126 ns/iter (+/- 19,085,347)
test bench_10000_trial_division   ... bench:   2,958,690 ns/iter (+/- 6,219,626)

とりあえず gmp のものは問題外に遅かった。次に遅いのが trial_division 。そして eratosthenes だと 10,000番目は 1.78 ms程度、で 100,000番目だと 30.7 ms程度。これだけでもPython版より20倍くらい速い…

そして atkin は eratosthenes と比較しても2倍くらい速い！すごい！！

高速化のテクニック

エラトステネスの篩よりアトキンの篩の方が速いのでそれを使おう、で終わりにしても良いかもしれないけど、エラトステネスの篩を使う場合でも色々工夫すれば速くしていける。

例を幾つか

上限個数を見積もる

篩のアルゴリズムの性質上、「limitまでの数を用意して フラグ管理していくことで素数を列挙していく」ということしか出来ず、その結果が何個になるかは最終出力を見てみないと分からない。

「N番目の素数を求めたい」というときに limitを幾つに設定して篩にかけていけばN番目までの素数を導き出せるかが不明なので、前述の実装だと limit = 1000 から始めて素数列を列挙し、N個に満たなければ limit を倍々にしていってN番目が求められるまで繰り返していっている。

10,000番目を求めるためには limitが128000になるまで8回、100,000番目を求めるためにはlimitが2048000になるまで12回、list_primesを呼び出して毎回同じような篩の操作をしていることになる。

この繰り返しも無駄だし、limitが無駄に大きすぎても N番目以降の数まで篩にかける操作が発生して無駄になる。

これを避けるために、N個の素数を返すギリギリのlimitの値を設定してあげたい。 ある自然数までに含まれる素数の個数を求める研究はたくさんされているようで

素数定理 - Wikipedia

素数計数関数 - Wikipedia

などで既に求められているものを使うとかなり近いものが出せそう。とりあえずは足りなくならない程度に雑に多めに見積もってやってみる。

pub fn eratosthenes_pi(n: usize) -> u32 {
    let n_ = n as f64;
    let lg = n_.ln();
    let limit = std::cmp::max(100, (n_ * lg * 1.2) as usize);

    let mut primes = Vec::new();
    let mut is_prime = vec![true; limit + 1];
    is_prime[0] = false;
    is_prime[1] = false;
    for p in 0..=limit {
        if !is_prime[p] {
            continue;
        }
        primes.push(p as u32);
        if primes.len() == n {
            return primes[n - 1];
        }
        for i in (p * p..=limit).step_by(p) {
            is_prime[i] = false;
        }
    }
    unreachable!();
}

これで従来のeratosthenesなどと比較してみると

test bench_100000_atkin           ... bench:  13,750,106 ns/iter (+/- 5,263,586)
test bench_100000_eratosthenes    ... bench:  30,559,236 ns/iter (+/- 7,994,169)
test bench_100000_eratosthenes_pi ... bench:  10,841,103 ns/iter (+/- 7,613,241)
test bench_10000_atkin            ... bench:     984,568 ns/iter (+/- 331,771)
test bench_10000_eratosthenes     ... bench:   2,210,553 ns/iter (+/- 2,621,658)
test bench_10000_eratosthenes_pi  ... bench:     907,250 ns/iter (+/- 254,367)

これだけで格段に速くなり atkin よりも高速になった。(勿論atkinでも同様の最適化すればもっと速くなるだろうけど)

Wheel factorization

そもそもlimitまでの数を調べていくのに半分は偶数で明らかに素数じゃないし、3の倍数のものも33%ほど含まれていて無駄。5の倍数だってそれなりに存在している… ということで無駄なものを最初から省いて調べていくのが良さそう。

ということを考えると、{2, 3, 5}から始めると そこから続く 7, 11, 13, 17, 19, 23, 29, 31 だけが2の倍数でも3の倍数でも5の倍数でもなく、そこから先は30ずつの周期で同様の間隔で増加させていった数値だけ見ていけば良い。 増加周期は [4, 2, 4, 2, 4, 6, 2, 6] となり、37, 41, 43, 47, 49, 53, 59, 61, ... と増えていく。もちろんこれは2, 3, 5しか見ていないので 7の倍数である49などは残る。これらをエラトステネスの篩で消していけば良い。

こういう手法を Wheel factorization と呼ぶらしい。

Wheel factorization - Wikipedia

とりあえずは単純に for p in 0..=limit で1つずつ順番に見ていたloopの部分だけ変更。

pub fn eratosthenes_wf(n: usize) -> u32 {
    let n_ = n as f64;
    let lg = n_.ln();

    let limit = std::cmp::max(100, (n_ * lg * 1.2) as usize);

    let mut primes = vec![2, 3, 5];
    let mut is_prime = vec![true; limit + 1];
    is_prime[0] = false;
    is_prime[1] = false;

    let inc = [6, 4, 2, 4, 2, 4, 6, 2];
    let mut p = 1;
    for i in 0.. {
        p += inc[i & 7];
        if p >= limit {
            break;
        }
        if !is_prime[p] {
            continue;
        }
        primes.push(p as u32);
        if primes.len() >= n {
            return primes[n - 1];
        }
        for j in (p * p..=limit).step_by(p) {
            is_prime[j] = false;
        }
    }
    unreachable!();
}

Benchmark結果は…

test bench_100000_atkin           ... bench:  25,911,095 ns/iter (+/- 46,670,614)
test bench_100000_eratosthenes    ... bench:  33,172,283 ns/iter (+/- 24,657,454)
test bench_100000_eratosthenes_pi ... bench:  11,062,096 ns/iter (+/- 4,717,035)
test bench_100000_eratosthenes_wf ... bench:   5,971,694 ns/iter (+/- 3,127,972)
test bench_10000_atkin            ... bench:     936,174 ns/iter (+/- 178,170)
test bench_10000_eratosthenes     ... bench:   1,790,384 ns/iter (+/- 711,067)
test bench_10000_eratosthenes_pi  ... bench:     797,356 ns/iter (+/- 171,738)
test bench_10000_eratosthenes_wf  ... bench:     399,302 ns/iter (+/- 48,778)

これだけでeratosthenes_piよりさらに2倍以上速くなった！単純なeratosthenesと比較するともう圧倒的ですね。

オチ

とこのように、エラトステネスの篩を使う手法にも様々な最適化の手段があり、それらを盛り込んでいると思われる primal というcrateがあります。

https://crates.io/crates/primal

これを使って StreamingSieve::nth_prime() でN番目の素数を求められる。

pub fn primal(n: usize) -> u32 {
    primal::StreamingSieve::nth_prime(n) as u32
}

中では 入力値の範囲によってより近似された (n番目の素数を含む下限/上限値)を見積もっていたり、ビット演算を駆使して高速に篩をかけるようにしているようだ。

Benchmark結果は…

test bench_100000_atkin           ... bench:  25,911,095 ns/iter (+/- 46,670,614)
test bench_100000_eratosthenes    ... bench:  33,172,283 ns/iter (+/- 24,657,454)
test bench_100000_eratosthenes_pi ... bench:  11,062,096 ns/iter (+/- 4,717,035)
test bench_100000_eratosthenes_wf ... bench:   5,971,694 ns/iter (+/- 3,127,972)
test bench_100000_primal          ... bench:     134,676 ns/iter (+/- 11,903)
test bench_10000_atkin            ... bench:     936,174 ns/iter (+/- 178,170)
test bench_10000_eratosthenes     ... bench:   1,790,384 ns/iter (+/- 711,067)
test bench_10000_eratosthenes_pi  ... bench:     797,356 ns/iter (+/- 171,738)
test bench_10000_eratosthenes_wf  ... bench:     399,302 ns/iter (+/- 48,778)
test bench_10000_primal           ... bench:       9,083 ns/iter (+/- 3,915)

はい、さらに数十倍速くなっていて完全に大勝利です。N番目の素数を求めたければこれを使いましょう。

Repository

• GitHub - sugyan/nth-prime-benchmark

ツッコミあればご指摘いただけると助かります。

References

• 素数列挙について - MugiCha
• Rubyで、「エラトステネスのふるい」を改良したと言われる「アトキンのふるい」をいまさら実装する - Line 1: Error: Invalid Blog('by Esehara' )
• 素数候補を効率よく絞りこむ「Wheel factorization」について調べたこと - Line 1: Error: Invalid Blog('by Esehara' )
• Sieve of Atkin - GeeksforGeeks
• Wheel Factorization Algorithm - GeeksforGeeks
• エラトステネスの篩の高速化 - Qiita

Advent of Code 2020 完答した

昨年はじめて真面目に挑戦した、Advent of Code。

memo.sugyan.com

今年もやるぞー！と意気込んで、12月になる前からちょいちょいTwitterで宣伝したりした甲斐もあってか、今年は日本国内でも挑戦している人が少し増えたようだった。嬉しい。

僕は昨年に引き続きRustで挑戦した。

github.com

出来る限り自力で考えて解いてみて、どうしても無理だったら他の人の答えを見る、くらいのつもりではいたけど、一応ぜんぶ自力での正解は出来た、と思う。たぶん。 (ひどい方法でとりあえず答えだけ出して、あとで他の人の答えを見て直した、というのも沢山あった。)

やったぁぁぁぁぁ！ / I just completed all 25 days of Advent of Code 2020! https://t.co/GnQF0hSuhM #AdventOfCode

— すぎゃーん💯 (@sugyan) December 25, 2020

問題の傾向としては昨年と比較すると今年は少し易しめだったかな？ 入力のparseがとにかく面倒、みたいな問題が多かった気がする。あと昨年のIntcodeシリーズのような「他の日の問題で実装したものを使って解く」系は無かった。それが良かったのかどうかは分からないけど (正直ちょっと期待していた)。 あと今年はJohn Horton Conway氏が新型コロナウイルスにより亡くなったというのもあってか、ライフゲーム(Conway's Game of Life)的なテーマのものが多かったようだ。

数学的に難しかったのは day13 だっただろうか。実装がとにかく大変だったのが day20。計算量が多くて厳しいのが幾つかあった。個人的には day23 がなんかすごく悩んでしまって一番躓いた。

せっかくなので後から遡ってロジックやコードを整理しつつ、自分なりの解答例と考え方をZennに書き残していっている。

sugyanの記事一覧 | Zenn

全部書き上げたら 2019 のも解き直してどんどん過去問も遡ってやってみようかな…

顔画像生成のためのデータセットを厳選する

memo.sugyan.com

の記事を書いてから10ヶ月も経っていた…。

• Twitterからアイドルの自撮り画像をひたすら収集
• dlib で顔検出し、各部位座標を利用して正規化し切り抜き
• Webアプリで管理、選別作業
• 選別作業自体を分類モデルに学習させて半自動化

というところの仕組みまでは出来て、あとは淡々と作業していくだけ、だったが まぁ色々あって時間かかり、、ここ最近でようやく選別後の画像が 16,500件ほど集まった。

十分とは言えないがそれなりの量にはなったので生成モデルを使って画像生成を試したい、と思ったが、改めて選別した画像を見返してみると「これは OK ラベルを付けたが やっぱり NG かな…」というものが幾つか混ざっていることに気付いた。

annotationの基準が変わるのはよくあることなので、初期は OK と思っていたものでも より良質な画像を多く見た後だと NG に分類すべきという判断になったりする。

とはいえ 16,500件の画像を一つ一つまたチェックするのも大変なので、自動的にある程度フィルタリングできるようにしていった。

重複検出

まずは重複している画像を排除。画像収集はTwitterからだけど、アイドルさんも同じ自撮り画像(もしくは僅かに加工をしただけのもの)を複数回使い回したりすることがある。選別時点では OK か NG かだけでしか判断していないので、後から見返してみると幾つもある同じ画像に OK ラベルをつけていることに気付いたりする。

わずかに加工されていたりしているとまったく同一の画像にはならないので、ここでは Perceptual Hash で類似しているものを抽出するようにした。 Python の ImageHash moduleで計算できる。

import imagehash
from PIL import Image

image = Image.open(image_path)
phash = imagehash.phash(image)

全候補画像に対しこれを計算し、値が同一もしくは限りなく近い値になっているものは重複として検出することができる。

顔検出を再度

収集した画像はすべて前記事の手法で dlib を使ってface detectionとlandmark detectionに成功して抽出されているもののはずだが、100%正確に検出できているとは限らない。元画像からの回転・切り抜き・拡大縮小などの影響で detection の結果が大きく変わっていたりするかもしれない。ので改めて前候補画像に対し検出をかけて face landmark の座標を算出し直した。

単一の顔だけが写っているはずなのに 実際に dlib.get_frontal_face_detector() にかけてみると一つも検出されなかったり複数検出されてしまったりする。このdetectorは run methodで幾つかのパラメータを指定して検出精度を調整できるらしいので、少しずつ調整して1つだけ検出されるようにした。

http://dlib.net/python/index.html#dlib.fhog_object_detector.run

import dlib

def detect_single_face(detector, image):
    for adjust_threshold in [x / 20.0 for x in range(0, -10, -1)]:
        for upsample_num_times in range(0, 3):
            detections, scores, indices = detector.run(
                image, upsample_num_times, adjust_threshold
            )
            if len(detections) == 1:
                return detections, scores, indices

正しく検出され正規化されている顔画像であれば、目の位置は平行になっていて顔の部分は中心で同程度の幅・高さを持っているはず。 それらの角度や座標の値が全候補画像の中で外れ値になっている場合、正しく顔やlandmarkを検出できておらず、学習データとして相応しくない可能性があるので除外することにした。

偏差を不偏標準偏差で割った検定統計量 が大きく平均値からブレているものを抽出。

ja.wikipedia.org

実際いくつかチェックしてみると大きく傾いてしまっているものやlandmarkが全然正しくない位置で検出されて正規化が失敗しているものなどが発見された。

エッジ強度でぼやけ抽出

次に、「正しく顔は検出できているが そもそも顔がハッキリ写っていないもの」を除外。 ピントが合っていないものやブレているものや暗くて画質が低いもの、など色んなパターンがあるので一概には決められないが、ともかく顔部分がハッキリ写っていないのは良くないので、候補から外していく。

ポートレートなどでは人物はハッキリ写っていて背景がぼやけていたりするので、まずは画像から顔領域だけを切り取って、その領域の画像に対して cv2.Laplacian を算出し、それらの値の分散を調べる。これが低いものは全体的に勾配が小さくハッキリと写っていないことが多い、と判断できる。

import cv2

img = cv2.imread(image_path, cv2.IMREAD_COLOR)
face = img[face_top:face_bottom, face_left:face_right, :]
gray = cv2.cvtColor(face, cv2.COLOR_BGR2GRAY)
variance = cv2.Laplacian(gray, cv2.CV_64F).var()

もちろん例外もあるだろうし、「この値より低かったらぼやけている」という閾値も決めづらいところではある。 今回は全体の中から下位1%を除外することにした。

口の有無

また、モノを食べていたりマスクを着けていたりして口がハッキリ写っていない画像が幾つかあったのが気になった。 そういったものが含まれても良いとは思うけど、出来れば今回のデータセットでは口がちゃんと写っている画像で揃えたい。

検出したlandmarkの座標を使えば、口の部分だけを切り取ることはできる。その領域の画像内に、「赤っぽい箇所」が発見できなかったら口が隠れていると判定することにした。 単純なRGBでのredの強さだけでは白色も赤とみなされてしまう。こういうときは HSV色空間で判定するのが良いらしい。Hueで赤色周辺の値で絞り、Saturation, Valueはそれなりの閾値を設けておく。その範囲内であれば「赤っぽい」といった判定ができるようだ。

import cv2
import numpy as np

mouth = img[ymin:ymax, xmin:xmax, :]
hsv = cv2.cvtColor(mouth, cv2.COLOR_BGR2HSV)
mask = sum(
    [
        cv2.inRange(hsv, np.array([0, 96, 128]), np.array([30, 255, 255])),
        cv2.inRange(hsv, np.array([150, 96, 128]), np.array([180, 255, 255])),
    ]
)
red = mask.sum()

赤は cv2 のHSVでは 0 - 180 の値で表されるらしい。ので、 0 - 30, 150 - 180 の範囲を赤とすることにした。S, V の範囲は適当に。 こうして得られた mask がすべて 0 だったなら、「口のあるはずの領域に赤い要素がまったく含まれていない」とみなすことができる。 実際幾つか調べてみると 食べ物などで口が隠れていたりするものが検出できた。それ以外にも普通に写っているはずのものもあったが、それらは画像全体が極端に明るかったり淡い色に加工されていたりして口の赤色が閾値に収まっていなかったようだ。 それはそれで学習データとしては微妙なので除外することにした。

全体の画質

顔がぼやけた画像などは極力排除したが、それ以外にも画像全体が画質が低いものもあったようだった。これも出来れば除外したい。 blind or No-reference (NR) image quality assessment (IQA) の手法としてBRISQUEというのがあるらしい。

image-quality moduleで計算できる。

from imquality import brisque
from PIL import Image

image = Image.open(image_path)
score = brisque.score(image)

これもどれくらいで閾値を設けるかは難しいところで、ともかく値が極端に高いものは画質が低いとみなすことが出来るようだったので、ある程度フィルタリングした中から 最終的にデータセットに選ぶときにはこの値が低い順に抽出するようにした。

最終結果

こうして、 OK ラベルをつけた 約16,500件の画像から 「重複なく、正しく顔検出された結果 大きな傾きや偏りなく顔領域が存在しており、ぼやけておらず、口も写っており、画質も低くない」ものを厳選して 16000枚の画像データセットを作成した。

そのデータを使って実際に StyleGAN2 ADA の生成モデルに学習させてみているところ それなりのアイドル顔画像が生成できてきているようだ。

最近のアイドル顔画像生成の状況 pic.twitter.com/ZJt2vpMjv2

— すぎゃーん💯 (@sugyan) November 17, 2020

この先についてはまた別記事に書こう

Repository

https://github.com/sugyan/image-dataset/tree/master/python/assessment

ISUCON10 予選敗退した

ISUCON10。 今年はあっという間に募集が終わって不参加かな、とも思ったけど 声かけていただき id:Soudai さんと id:kamipo さんと、昨年と同じチームで出場した。

• ISUCON9 予選敗退した - すぎゃーんメモ
• ISUCON9 本選12位だった - すぎゃーんメモ

とはいえ今年は僕は子も産まれ 京都に移住したのもあって、昨年のように東京で3人集まっての参加は難しく、一人だけ京都の自宅からリモートで繋いで参加、という形になった。

素振り

一応去年組んだメンバーであるとはいえ 敗退しているわけだし、今回は前述のように一人はリモートで作業することになるので練習はしておくべきだろう、ということで 2週間ほど前に1日つかって素振り(練習)をした。 昨年の予選問題はすっかり忘れているのでそれを使って、一日中zoomで繋いで会話しつつ、必要に応じてエディタの画面を共有して一緒に見ながら作業する練習だったり。 あと今年は New Relic が特別ライセンスで使用できるってことでせっかくだからそれを導入してみてどうなるか、の練習も。

当日

朝5時台に子が目覚めるので一緒に起きて抱っこで30〜40分ほど散歩し、帰宅してから朝食たべて 子の朝寝と一緒に再度寝て、9時頃目覚めてさぁ準備、と思ったら2時間の延期…。 18時にはお風呂に入れて20時までには寝かしつけ、というサイクルだったのでどうしよう…という困惑はありつつも、どうしようもないので今日は妻に頑張っていただく、ということで ともかく開始まで待機。

序盤

まずは環境作り。サーバにログインしてアプリの動作を確認、git repositoryに最低限のコードをpushして ローカル環境で動作できるかどうか試す。 web app がAPIのみで static files が分離されていたのでちょっと苦労したけど、それらもローカルに持ってきて手元でnginxも立てたりして とりあえず動かせるようにはした。少し手間はかかったけど これによってコード変更時の動作確認とデバッグがやりやすくなっていたのでやっておいて正解だったと思う。

Rubyでやっていく方針だったので初期実装のGoからRubyに切り替え。NewRelicを仕込んでベンチを回してボトルネックを見てみる。とにかく検索が重い。

インフラ周りやDBのチューニングはSoudaiさんに任せてアプリやクエリの改善をしていこう、と。 まずはじゃあ明らかにヤバそうな features の LIKEクエリを別tableに正規化して軽くしよう、と動き始め。

中盤

estate_features, chair_features のtableを作って 初期データから生成されるデータを作って insert時にfeaturesもinsertしていく処理を書いて 検索時にそれらを使ったクエリに書き換えて…

OR検索じゃなくてAND検索だったから GROUP BY にしなきゃ、id IN (?) だと対象idsがemptyになった場合はクエリ壊れてしまう、とかでちょいちょいハマったり。 結局これがちゃんと動いたのは16:30過ぎ。時間かかりすぎたな…。しかも実際には features を使ったクエリは少なかったようであまりスコアには影響せず (ここは反省点)。

じゃあまだ遅いのは何だろうね、ということで 幾つかindexを適当に貼ったりしながら次の改善点を探す。 width や height などのrangeで絞っているのを 選択肢idのequalで絞れないか、と提案。やってみることに。 kamipoさんにschema変更してもらって そこの初期値を設定するコードを書いてもらい、それを使ってクエリを単純化するコードを書いて、と分担し。

しかしこのあたりで POST /api/estate や POST /api/chair がタイムアウトしてしまうなどでベンチがまったく通らず。 諸々改善してようやく通るようになった頃にはもう20時を過ぎていた。これでスコアはようやく 1,500 程度。。

GET /api/recommended_estate/:id はSoudaiさんが UNION クエリに書き換えたが 何故かベンチがこけるのでrevert。 やっていることをよく見て考えたら 「椅子の3辺のうち2つが ドアの2辺におさまるかどうか」だけを調べられれば良いのではないか、と気付いたので少しだけ簡略化することは出来た。このへんは普段leetcodeとかで問題読んで目的のコードに落とし込む訓練をしていた成果はあったかな、と思う。

終盤

やり残したことはたくさんあって nazotte とかはほぼ手を付けられず (kamipoさんが最低限の無駄なクエリを省くようにはしてくれていたが)、あとはもう時間も無いので logを切ったり NewRelicを切ったり 再起動テストをしたり で終了。最終スコアは 1,656 と 予選通過ボーダーにも全然届かない結果となった。

反省点

自分なりにはベストを尽くして最大限のパフォーマンスは出せたと思っているので悔いは無いけど、それでこの結果なので単純に力不足が露呈することになっていてつらい気持ち。

チームの反省として、kamipoさんの力を引き出せなかった、というのがある。予選翌日に感想戦で改めて見直したところ MySQLのindexまわりで改善できるところがたくさんあった。

• Pull requests · soudai/isucon10-qualify · GitHub

このあたりに自分ももう少し気付ければ良かったのは勿論だけど、チームとして最高のパフォーマンスを出すためには もっとkamipoさんに集中してここらへんを見てもらえるように作業分担を調整する必要があったのだろうな…。

features の正規化 (時間かかったわりにはそこまで大きなスコア改善に寄与しなかったわけだが)や rangeの選択肢のequal化などは優先度を調整しつつ僕が一人で全部やって その間にkamipoさんがindexを見れるように、といった分担が出来ていなければいけなかったんだろうな、と思う。

昨年もそうだったけど全体的に手を動かすスピードが遅かったりハマったときに無駄にハマり続けて時間を浪費してしまっていた、というのはある。 ここは日々の鍛錬でどうにかしていくしかない。

感想

一人自宅からリモート参加というのは初挑戦だったけど、そこはそれほど苦にはならない程度に出来たかな、と思うので それは良かったと思う。 とはいえ理想としてはチーム全員集まってワイワイやれた方が楽しいとは思うけども。 来年はどうなるかな〜。

運営の皆様は今回も準備や当日の進行が大変そうではありましたが 取り組み甲斐のある面白い問題を出していただき感謝です。 おつかれさまでした。本選も素晴らしい競技になることを期待しております。

Repository

• https://github.com/soudai/isucon10-qualify

StyleGAN2学習済みモデルを使って任意の画像を生成させる

memo.sugyan.com

の続き。

StyleGAN2 は "mapping network" と "synthesis network" の2つのネットワークで構築されていて、画像の生成を行う synthesis network への入力 dlatents_in を変化させていくことで様々な変化を出せる、というものだった。

前回は mapping network からの出力値を使って 「学習によって上手く生成できるようになった画像」のための dlatents_in の値の間を遷移させるといったことをしていたけど、実際には synthesis network には十分に様々な画像を生成できる能力が獲得されているはず、らしい。 具体的には、アイドルの顔画像だけで学習したモデルでも アイドルの顔以外の画像も生成できるかもしれない、ということ。

以下の論文で様々な実験・検証が行われている。

• Image2StyleGAN: How to Embed Images Into the StyleGAN Latent Space?

任意画像を生成するための latent space の学習

要するに synthesis network への入力 dlatents_in を (14, 512) の shape を持つ変数 (14 は 256x256サイズの場合の数値) とみなし、それを使って生成される画像が目標画像に近くなるように学習させていけば良い、ということ。

せっかくなので TensorFlow 2.x で動くように書いてみた。

Snapshot から SavedModel への変換

StyleGAN2 の公式実装は TensorFlow 1.x でしか動かない。ので 以前の記事 に書いたように snapshot の .pkl から Generatorの部分だけ取り出して SavedModel の形式に変換して保存する。

output_names = [t.name for t in Gs.output_templates]

with tf.Graph().as_default() as graph:
    outputs = [graph.get_tensor_by_name(name) for name in output_names]
    images = tf.transpose(outputs[0], [0, 2, 3, 1])
    images = tf.saturate_cast((images + 1.0) * 127.5, tf.uint8)
    # save as SavedModel
    builder = tf.compat.v1.saved_model.Builder(save_dir)
    signature_def_map = {
        'synthesis': tf.compat.v1.saved_model.build_signature_def(
            {'dlatents': tf.saved_model.utils.build_tensor_info(outputs[1])},
            {'images': tf.saved_model.utils.build_tensor_info(images),
             'outputs': tf.saved_model.utils.build_tensor_info(tf.transpose(outputs[0], [0, 2, 3, 1]))})
    }
    builder.add_meta_graph_and_variables(
        sess,
        [tf.saved_model.tag_constants.SERVING],
        signature_def_map)
    builder.save()

最終的な画像としての出力は [0, 255] の tf.uint8 の値に変換したものだけど、ここではその前段階での synthesis network の出力を NHWC に変換しただけのものを使う。 これは (?, 256, 256, 3) の tf.float32 tensor で、 [-1.0, 1.0] の範囲の値を持つとみなして処理される。

Keras layers の構築

TensorFlow 2.x では主に Keras API を使って model の構築 & 学習 をしていくことになる。 tf.keras.layers.Layer を継承した独自の layer を定義していく。

class LatentSpace(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__(input_shape=())
        self.v = self.add_weight(
            shape=(1, 14, 512),
            dtype=tf.float32)

    def call(self, inputs):
        return tf.identity(self.v)


class Synthesis(tf.keras.layers.Layer):
    def __init__(self, model_path):
        super().__init__()
        model = tf.saved_model.load(model_path)
        self.synthesis = model.signatures['synthesis']

    def call(self, inputs):
        return self.synthesis(dlatents=inputs)['outputs']

まずは (1, 14, 512) の変数だけを持つ layer。 入力されてくる inputs は無視して、持っている変数をそのまま出力する。この add_weight で登録された変数たちが、training すべきパラメータとなる。

次に その変数の値を入力として受けて 生成を行う layer。 これは先述した SavedModel を load して結果を返してやるだけで良い。

これで Model の構築ができる。

model = tf.keras.Sequential([
    LatentSpace(),
    Synthesis(model_path),
])
model.summary()

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #
=================================================================
latent_space (LatentSpace)   (1, 14, 512)              7168
_________________________________________________________________
synthesis (Synthesis)        (1, 256, 256, 3)          0
=================================================================
Total params: 7,168
Trainable params: 7,168
Non-trainable params: 0
_________________________________________________________________

Target image と Dataset

生成目標となる画像を用意する。普通に読み込んで decode すると [0, 255] 範囲の tf.uint8 tensor になってしまうので、synthesis network の出力に合わせて [-1.0, 1.0] の範囲になるよう調整する。

それを使って 学習時の入力データを作成する。といっても Model への入力は不要なので適当に 0 とかを返しておく。 Target data となる y だけ常に同じ値を返し続ければ良い。

with open(target_image, 'rb') as fp:
    y = tf.image.decode_jpeg(fp.read())
y = tf.expand_dims(tf.cast(y, tf.float32) / 127.5 - 1.0, axis=0)

dataset = tf.data.Dataset.from_tensors((0, y))

Loss class

あとは最小化すべき loss の定義。 論文によると 生成画像と目標画像の間の pixel-wise MSE と、 VGG16 を使った perceptual loss を組み合わせて使うようだ。 要するに pixel 間の差分だけでなく 特徴も似たようなものになるのが良い、ということのようで。

tf.keras.losses.Loss を継承した独自の loss を定義する。

class EmbeddingLoss(tf.keras.losses.Loss):
    def __init__(self, image):
        super().__init__()
        self.vgg16 = tf.keras.applications.VGG16(include_top=False)
        self.target_layers = {'block1_conv1', 'block1_conv2', 'block3_conv2', 'block4_conv2'}
        self.outputs = []
        out = image
        for layer in self.vgg16.layers:
            out = layer(out)
            if layer.name in self.target_layers:
                self.outputs.append(out)

    def call(self, y_true, y_pred):
        out = y_pred
        outputs = []
        for layer in self.vgg16.layers:
            out = layer(out)
            if layer.name in self.target_layers:
                outputs.append(out)
        n = tf.cast(tf.math.reduce_prod(y_pred.shape), tf.float32)
        losses = tf.math.reduce_sum(tf.math.squared_difference(y_true, y_pred)) / n
        for i, out in enumerate(outputs):
            n = tf.cast(tf.math.reduce_prod(out.shape), tf.float32)
            losses += tf.math.reduce_sum(tf.math.squared_difference(self.outputs[i], out)) / n
        return losses

VGG16 の model は tf.keras.applications に imagenet で学習したものがあるようなので それをそのまま使う。論文によると conv1_1, conv1_2, conv3_2, conv4_2 の4つの layer の出力を使ってそれぞれ差分を足し合わせて loss の値にしている、とのこと。 目標画像については値が変化しないので、この中間層の特徴量も変化しない。ので最初に計算して保持しておく。 call() 時には y_pred で model からの出力値が渡されてくるので、その都度 VGG16 に通して各層の出力を取得する。 それぞれ目標値との tf.math.squared_difference を tf.math.reduce_sum して、それぞれの scale で割ってやる。 最終的な和が、最小化すべき loss の値になる。

学習

ここまで出来たらあとは compile して fit させるだけ。 前述の EmbeddingLoss を loss に指定して、 Adam optimizer で最適化していく。 実験してみた感じではこの optimizer のパラメータによって学習の結果も大きく変わってくるようで、このへんの最適な値を見つけるのはとても難しそうだった。 ここでは論文記載と同じ learning_rate=0.01, epsilon=1e-08 を使用する。

fit では 前述の dataset を batch_size: 1 で繰り返す。適当な stepsで 1 epoch の区切りにして、その epoch 終了時の変数での生成結果を画像として出力するようにする。

class GenerateCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs):
        v = self.model.layers[0].variables[0].numpy()
        images = self.model.layers[1](v)
        images = tf.saturate_cast((images + 1.0) * 127.5, tf.uint8)
        with open(f'epoch{epoch:03d}.png', 'wb') as fp:
            data = tf.image.encode_png(tf.squeeze(images, axis=0)).numpy()
            fp.write(data)


model.compile(
    optimizer=tf.keras.optimizers.Adam(
        learning_rate=0.01,
        epsilon=1e-08),
    loss=EmbeddingLoss(y))
model.fit(
    dataset.repeat().batch(1),
    steps_per_epoch=50,
    epochs=100,
    callbacks=[GenerateCallback()])

さすがにこれは CPU環境ではそれなりに時間がかかって厳しい。 Google Colaboratory の GPU Runtime だと数分で 5,000 stepくらいは完了するようだ。

学習結果

自力で収集した アイドルの顔画像 7,500 枚で ある程度学習したモデルを使用。

まずは 実際にこのモデルによって生成された画像。

これは既に自らが生成した実績がある画像なので、再現できないとおかしいくらいのものではある。

意外と「完全に一致」というところまではいかない…。けど まぁ早い段階からほぼ再現できているようには見える。 ここがより安定して近くなるかどうかは optimizer のパラメータ次第という感じではあった。

次に、この生成モデルへの学習にはまったく使っていない 女優さんの顔画像とかだとどうなるだろうか。

ちょっとハッキリしないけど、一応それなりに髪や顔のパーツまで生成できているようだ。

では もはや日本人でも若い女性でもない人物の顔画像を目標にした場合は…？

思ったよりイケる！ これはこれで予想外。

まったく学習データに使っていない画像も生成できるようになる、というのは面白いな〜。 今回使ったモデルはかなり学習データのドメインが限定的だし 学習も完了ってほど十分に出来ていないのだけど、それでもこれだけ生成できる、ということが分かった。

もっと学習が進んだものや データを増やして学習させたモデルを使った場合にはまた違う結果になるかもしれない。

おまけ: morphing

こうして任意の画像を生成するための synthesis network への入力が得られた、ということは 前回の記事 で書いたように、2つの入力があった場合にその間の値を使うことで morphingが出来る…はず。

と思ってやってみたが

と 中間の表現は気持ち悪いばかりのものになってしまった。

これくらい離れた空間同士だと単純な線形の推移では自然な変化を出せないようだ。このあたりも学習の進行度合いによって違ったりするかもしれないけど…。

Repository

• https://github.com/sugyan/gan-playground

StyleGAN2学習済みモデルを使ったmorphing、latent spaceの探求

学習データはまだまだ収集途中だし 学習もまだ完了とは言えない状態なのだけど、なんとなくそれっぽい顔画像は生成できるくらいまでは出来てきているので、それを使った実験をしてみている。

学習した生成モデルを使って morphingを試してみてる #StyleGAN2 pic.twitter.com/NxlaOgp1RX

— すぎゃーん💯 (@sugyan) February 7, 2020

なるほどねー pic.twitter.com/7Qmf61G14t

— すぎゃーん💯 (@sugyan) February 9, 2020

主にこの記事を参考にしています。

qiita.com

1. latents_in の線形移動

まず最初に試したのは、通常の generator network を使った場合の latents_inを使ったもの。

generator は、(?, 512) の任意の入力を受けて画像を生成するようになっている。 これが所謂 潜在空間 (latent space) というやつで ここの値をうまく選ぶことで所望の出力を得られるようになったりする、というわけだ。

例えば適当な乱数で (2, 512) を作ってそれを入力すると、2つの異なる画像が出力される。

その2つの (1, 512) の乱数ベクトルの、中間の値を入力として使えば 生成される画像も2つの画像の中間のものになるだろう、という考え方。

import numpy as np
import tensorflow as tf
from PIL import Image

model_path = '...'
model = tf.saved_model.load(model_path)
generator = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

rnd = np.random.RandomState(0)
z = rnd.randn(2, 512)

inputs = []
steps = 10
for i in range(10):
    inputs.append(z[0] + (z[1] - z[0]) * i / steps)

for i, latents in enumerate(inputs):
    images = generator(latents=tf.constant([latents], tf.float32))['images']
    Image.fromarray(images.numpy()[0]).save(f'out_{i:02d}.png')

z[0] から z[1] までを 0.1 刻みで線形に変化させてそれぞれ画像を生成した結果。

右端と左端は(一応)違う顔だが、その間で徐々に変化していっているのが分かる、と思う。 GIFアニメーションにするとこんな感じ。

2. mapping networkの出力 dlatents_in を使う

StyleGAN の generator は、"mapping network" と "synthesis network" の2つの network によって作られている。 generator に入力された latents_in の (?, 512) の値は mapping network で (?, 14, 512) といった もう少し次元の大きい変数に変換される。(14 というのは 256x256 のときの数値で upsample の layer 数が増えるとまた 16 や 18 に変化したりする、のかな。) で、この Disentangled latents と呼ばれる出力が、後段の synthesis network への入力に使われ、実際の画像の合成が行われる、ということになる。 dlatents_in と呼ばれるこの値が 画像生成のための入力としてはより直接的なものになるのかもしれない。

ザックリした理解では、 synthesis network は実に広範囲な生成の能力を持っているが より所望の(学習データに近い)主力を得るための dlatents_in を作り出すための前段として mapping network が作用している、という感じか。実際には その出力からさらに truncation_psi trick といった より質の良い出力を得ることが出来る値にするための工夫がされているようだ。

ともかく、この 2つの network を分けて考えると、2つの画像間を補完する場合も mapping network への入力 latents_in を線形に変化させるよりも synthesis network への入力 dlatents_in を変化させた方が自然なものになりそう。

generator も2つに分けて各 network の入出力を扱えるようにした。 StyleGAN2 では return_dlatents = True と option を指定することで dlatents_in の tensor も得ることが出来る。

この dlatents_in をまず先に計算し、先程と同様にその値を 0.1 刻みで線形に変化させるようにする。

import numpy as np
import tensorflow as tf
from PIL import Image

model_path = '...'
model = tf.saved_model.load(model_path)
mapping = model.signatures['mapping']

rnd = np.random.RandomState(0)
z = rnd.randn(2, 512)
z = mapping(latents=tf.constant(z, tf.float32))['dlatents'].numpy()

inputs = []
steps = 10
for i in range(10):
    inputs.append(z[0] + (z[1] - z[0]) * i / steps)

synthesis = model.signatures['synthesis']
for i, latents in enumerate(inputs):
    images = synthesis(dlatents=tf.constant([latents], tf.float32))['images']
    Image.fromarray(images.numpy()[0]).save(f'out_{i:02d}.png')

先述の latents_in の変化だと途中で顎のあたりに手のようなものが現れたりガチャガチャと忙しい変化になってしまっていたのに対し、今度はシームレスに2画像間を推移するようになった。

latent space を探る

dlatents_in は本当に広い空間で、mapping network を通さずに本当に random な値を使って synthesis network に入力すると まったく汚い出力になってしまう。

逆にここを上手く最適化してやることで 学習データにまったく存在していない画像も生成することが出来る可能性があるようだ。これについてはまたこれから実験していきたい。

dlatents_in の値を mapping network の出力に絞ることで それなりに学習データに近い出力が得られるようにはなるが、学習データの偏りや学習不足などのせいか どうしても印象が似たような顔ばかりになりやすい。

とはいえ、これも mapping network の出力として得られる dlatents_in が大きく異なっていればそれなりには異なる出力画像になるのではないか？

1000件ほどの random な latents_in を mapping network に入力し、得られた各 dlatents_in の間の距離を計算してみる。これが最も遠い組み合わせを選んで使ったら、出力画像も大きく異なるものになるだろうか？

rnd = np.random.RandomState(0)
z = rnd.randn(1000, 512)
z = mapping(latents=tf.constant(z, tf.float32))['dlatents'].numpy()

distances = []
for i in range(z.shape[0]):
    for j in range(i + 1, z.shape[0]):
        distances.append([np.linalg.norm(z[i] - z[j]), (i, j)])

_, (i, j) = sorted(distances, reverse=False)[0]
z = [z[i], z[j]]

inputs = []
steps = 10
for i in range(10):
    inputs.append(z[0] + (z[1] - z[0]) * i / steps)

背景がゴチャゴチャしたり 画像の質はあまり良くないかもしれないけど、それなりにインパクトのある変化をする morphing が出来たような気がする。

ちなみに最も距離が近い2つを選択すると

となり、ほぼ違いが感じられないような変化になったので この感覚はまぁ合ってそうだ。

Repository

• https://github.com/sugyan/gan-playground