N番目の素数を求める

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

環境

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

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

エラトステネスの篩の計算量 O(N \log \log N) に対し こちらは O(\frac{N}{\log \log N}) になる、とのこと。

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

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

の3つの式において xy の組み合わせの数が合計で 奇数個存在している 場合に、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倍くらい速い…

そして atkineratosthenes と比較しても2倍くらい速い!すごい!!

高速化のテクニック

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

例を幾つか

上限個数を見積もる

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

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

10,000番目を求めるためには limit128000になるまで8回、100,000番目を求めるためにはlimit2048000になるまで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
}

中では 入力値の範囲によってより近似された p(n) (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

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

References

Advent of Code 2020 完答した

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

memo.sugyan.com

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

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

github.com

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

問題の傾向としては昨年と比較すると今年は少し易しめだったかな? 入力の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からだけど、アイドルさんも同じ自撮り画像(もしくは僅かに加工をしただけのもの)を複数回使い回したりすることがある。選別時点では OKNG かだけでしか判断していないので、後から見返してみると幾つもある同じ画像に OK ラベルをつけていることに気付いたりする。

わずかに加工されていたりしているとまったく同一の画像にはならないので、ここでは Perceptual Hash で類似しているものを抽出するようにした。 PythonImageHash 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を検出できておらず、学習データとして相応しくない可能性があるので除外することにした。

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

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()

赤は cv2HSVでは 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 の生成モデルに学習させてみているところ それなりのアイドル顔画像が生成できてきているようだ。

f:id:sugyan:20201120235412p:plain

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

Repository

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

ISUCON10 予選敗退した

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

とはいえ今年は僕は子も産まれ 京都に移住したのもあって、昨年のように東京で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を適当に貼ったりしながら次の改善点を探す。 widthheight などのrangeで絞っているのを 選択肢idのequalで絞れないか、と提案。やってみることに。 kamipoさんにschema変更してもらって そこの初期値を設定するコードを書いてもらい、それを使ってクエリを単純化するコードを書いて、と分担し。

しかしこのあたりで POST /api/estatePOST /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まわりで改善できるところがたくさんあった。

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

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

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

感想

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

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

Repository

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

memo.sugyan.com

の続き。

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

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

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

任意画像を生成するための 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。 これは先述した SavedModelload して結果を返してやるだけで良い。

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

学習

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

fit では 前述の datasetbatch_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 枚で ある程度学習したモデルを使用。

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

f:id:sugyan:20200216220601p:plain

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

f:id:sugyan:20200216220059g:plain

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

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

f:id:sugyan:20200216220015g:plain

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

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

f:id:sugyan:20200216220128g:plain

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

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

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

おまけ: morphing

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

と思ってやってみたが

f:id:sugyan:20200216215837g:plain

f:id:sugyan:20200216215904g:plain

f:id:sugyan:20200216215127g:plain

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

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

Repository

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

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

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

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 刻みで線形に変化させてそれぞれ画像を生成した結果。

f:id:sugyan:20200209192014p:plain

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

f:id:sugyan:20200209192737g:plain

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 数が増えるとまた 1618 に変化したりする、のかな。) で、この 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_intensor も得ることが出来る。

この 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')

f:id:sugyan:20200209200947p:plain

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

f:id:sugyan:20200209201028g:plain

latent space を探る

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

f:id:sugyan:20200209201952p:plain

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


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)

f:id:sugyan:20200209204147p:plain

f:id:sugyan:20200209204203g:plain

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

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

f:id:sugyan:20200209204427p:plain

f:id:sugyan:20200209204446g:plain

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

Repository

StyleGAN2による画像生成をCPU環境/TensorFlow.jsで動かす

memo.sugyan.com

の続き。

ようやくTensorFlow.jsを使ってブラウザ上で動かせるようになったので、そのためにやったことメモ。

(まだまだ画像の質とかパフォーマンスの問題とかは色々ある)

CPU環境で動かす

最終的にはTensorFlow.jsでブラウザ上で動かすことが目標だったので別にCPUで動かせる必要は無かったのだけど、どうもGPU環境でしか動かない特殊なOpなどもあってそれが変換後のモデルで実行時にエラーを引き起こす原因だったりするため、まずはCPU環境でも安定して動くようにするのが確実なようだ。

前回記事にも書いた通り、 StyleGAN2 の学習の過程で出力される .pkl ファイルは、CPU環境では読み込むことも出来ない。

import pickle
from dnnlib import tflib

tflib.init_tf()
with open('network.pkl', 'rb') as fp:
    _, _, Gs = pickle.load(fp, encoding='latin1')

Gs.print_layers()
...

RuntimeError: NVCC returned an error. See below for full command line and output log:

nvcc "/Users/sugyan/.ghq/github.com/NVlabs/stylegan2/dnnlib/tflib/ops/fused_bias_act.cu" ...

/bin/sh: nvcc: command not found

そりゃCUDAが入っていない環境ではこうなる。 ので、無理矢理にコードを書き換えてCUDAを使わない reference implementation を利用するようにする。 特に upfirdn_2d の方は幾つかの場所から呼ばれるので、デフォルト引数をまとめて書き換えるのも良いけど impl_dict を書き換えて reference implementation を強制してしまうのがラク

diff --git a/dnnlib/tflib/ops/fused_bias_act.py b/dnnlib/tflib/ops/fused_bias_act.py
index 52f6bfd..c294277 100755
--- a/dnnlib/tflib/ops/fused_bias_act.py
+++ b/dnnlib/tflib/ops/fused_bias_act.py
@@ -63,7 +63,7 @@ def fused_bias_act(x, b=None, axis=1, act='linear', alpha=None, gain=None, impl=

     impl_dict = {
         'ref':  _fused_bias_act_ref,
-        'cuda': _fused_bias_act_cuda,
+        'cuda': _fused_bias_act_ref,
     }
     return impl_dict[impl](x=x, b=b, axis=axis, act=act, alpha=alpha, gain=gain)

diff --git a/dnnlib/tflib/ops/upfirdn_2d.py b/dnnlib/tflib/ops/upfirdn_2d.py
index fd23777..1df2935 100755
--- a/dnnlib/tflib/ops/upfirdn_2d.py
+++ b/dnnlib/tflib/ops/upfirdn_2d.py
@@ -57,7 +57,7 @@ def upfirdn_2d(x, k, upx=1, upy=1, downx=1, downy=1, padx0=0, padx1=0, pady0=0,

     impl_dict = {
         'ref':  _upfirdn_2d_ref,
-        'cuda': _upfirdn_2d_cuda,
+        'cuda': _upfirdn_2d_ref,
     }
     return impl_dict[impl](x=x, k=k, upx=upx, upy=upy, downx=downx, downy=downy, padx0=padx0, padx1=padx1, pady0=pady0, pady1=pady1)

これで 引数が impl='cuda' で指定されてこようが関係なく ref の方が使われるようになる。

こうすると .pkl を読み込むのは問題なくできるようになる。 ただここから実際に計算を実行しようとすると問題が起きるわけで。

Gs.run()GPU device を必要とするので output tensor を指定して tf.Session.run() してみる。

import pickle
import numpy as np
import tensorflow as tf
from dnnlib import tflib

tflib.init_tf()
with open('network.pkl', 'rb') as fp:
    _, _, Gs = pickle.load(fp, encoding='latin1')

graph = tf.get_default_graph()
inputs = graph.get_tensor_by_name(f'Gs/{Gs.input_names[0]}:0')
outputs = graph.get_tensor_by_name(f'Gs/{Gs.output_names[0]}:0')

rnd = np.random.RandomState(0)
z = rnd.randn(1, *Gs.input_shape[1:])
print(tf.get_default_session().run(outputs, feed_dict={inputs: z}))
2020-02-04 23:26:46.471886: E tensorflow/core/common_runtime/executor.cc:642] Executor failed to create kernel. Invalid argument: Conv2DCustomBackpropInputOp only supports NHWC.
         [[{{node Gs/G_synthesis/8x8/Conv0_up/conv2d_transpose}}]]

...

StyleGAN2のモデルはGPU環境で学習する前提で作られているので、その環境に最適化された処理が幾つかある。そのうちの一つが NCHW のdata formatを使っていること。 この形でmodelが作られていると、CPU環境では NHWC にしか対応していないので計算を実行することが出来ないようだ。

1. Graphを書き換える

ということで困った、という記事を書いたところ、以下のようなフィードバックをいただいた。

なるほどー、構築されたGraphを舐めていって inputsoperation を書き換えることで NCHWNHWC に変換する方法があるのか…!

ということで上記を参考にしながら自分で書いてみた。

tflib.init_tf()
with open('network.pkl', 'rb') as fp:
    _, _, Gs = pickle.load(fp, encoding='latin1')

graph = tf.get_default_graph()

target_ops = []
for op in graph.get_operations():
    if not op.name.startswith('Gs/'):
        continue
    if 'data_format' in op.node_def.attr and op.node_def.attr['data_format'].s == b'NCHW':
        target_ops.append(op)

まずは Gs/ 以下の、 'NCHW' という data_format attribute を持つ operation をすべて抽出する。これらが書き換える対象となる。

for target_op in target_ops:
    print(f'op: {target_op.name} ({target_op.type})')
op: Gs/G_synthesis/4x4/Conv/Conv2D (Conv2D)
op: Gs/G_synthesis/4x4/ToRGB/Conv2D (Conv2D)
op: Gs/G_synthesis/8x8/Conv0_up/conv2d_transpose (Conv2DBackpropInput)
op: Gs/G_synthesis/8x8/Conv0_up/Conv2D (Conv2D)
op: Gs/G_synthesis/8x8/Conv1/Conv2D (Conv2D)
op: Gs/G_synthesis/8x8/Upsample/Conv2D (Conv2D)
op: Gs/G_synthesis/8x8/ToRGB/Conv2D (Conv2D)
op: Gs/G_synthesis/16x16/Conv0_up/conv2d_transpose (Conv2DBackpropInput)
op: Gs/G_synthesis/16x16/Conv0_up/Conv2D (Conv2D)

...

それらの operation に対し、まずは inputs のshapeを変換していく。

for target_op in target_ops:
    # Input tensors
    if target_op.type == 'Conv2D':
        inputs = [
            tf.transpose(
                target_op.inputs[0],
                [0, 2, 3, 1],
                name=f'{target_op.name}_input_transpose'),
            target_op.inputs[1]
        ]
    elif target_op.type == 'Conv2DBackpropInput':
        inputs = [
            tf.gather(
                target_op.inputs[0],
                [0, 2, 3, 1],
                name=f'{target_op.name}_output_shape_transpose'),
            target_op.inputs[1],
            tf.transpose(
                target_op.inputs[2],
                [0, 2, 3, 1],
                name=f'{target_op.name}_value_transpose')
        ]

実際に見てみると分かるが、こうして 'NCHW' な op を抽出してみると すべて typeConv2DConv2DBackpropInput のどちらかしかない。 少なくとも StyleGAN2 では、この2つのtypeに対してそれぞれ対応するだけで良い、ということになる。

Conv2D の場合、 inputs0 番目に入力のTensorが入ってくる。これが例えば (?, 1, 11, 11) だったりして (N, C, H, W) に対応する。 ので、 tf.transpose[0, 2, 3, 1] を指定することでこの入力Tensor(N, H, W, C) に変換することが出来る。 1 番目の入力は filter の値のようで、これはdata formatに依存しないのでこのまま使えば良い。

Conv2DBackpropInput の場合はもう少し厄介。どうやら入力は output_shape, filter, value という順番で来るらしい。 1 番目は Conv2D の場合と同様そのまま使って 2 番目が入力Tensorなので やはり同様に [0, 2, 3, 1] で transpose してやる。 そして 0 番目が その出力結果のshapeをどうするか指定するという役割のようで、そのshapeを示す (4,)Tensorが入ってくる。 これは例えば [1, 512, 9, 9] といった形の値で やはり NCHW ならその形に指定するわけだけど、ここではこの op を NHWC に変えてやりたいので、この output_shape も書き換えてやらないと その後の出力の型が合わなくなってしまう。 tf.gather を使ってこの output_shape の中身の順番を入れ替える。

これが出来たら次は attributes。

    # Attributes
    attrs = {}
    for k, v in target_op.node_def.attr.items():
        if k == 'data_format':
            continue
        if target_op.type == 'Conv2DBackpropInput' and k == 'strides':
            strides = v.list.i
            attrs[k] = tf.AttrValue(list=tf.AttrValue.ListValue(i=[
                strides[0],
                strides[2],
                strides[3],
                strides[1],
            ]))
        else:
            attrs[k] = v

各 operation は入力値とは別に? attributes というものを持っているようで、これも書き換えてやる必要がある。

data_formatNCHW である、という情報はここに含まれているので、変換する際にはこれを捨ててしまうことで defaultの NHWC にすることが出来る。

もう一つ Conv2DBackpropInput の場合に変更する必要があるのが strides の値で、これも NCHW のときと NHWC のときで扱いが変わるものらしい。

この strides と前述の output_shape については upfirdn_2d.upsample_conv_2d() に分岐が書かれている。

upfirdn_2d.py 抜粋:

    # Determine data dimensions.
    if data_format == 'NCHW':
        stride = [1, 1, factor, factor]
        output_shape = [_shape(x, 0), outC, (_shape(x, 2) - 1) * factor + convH, (_shape(x, 3) - 1) * factor + convW]
        num_groups = _shape(x, 1) // inC
    else:
        stride = [1, factor, factor, 1]
        output_shape = [_shape(x, 0), (_shape(x, 1) - 1) * factor + convH, (_shape(x, 2) - 1) * factor + convW, outC]
        num_groups = _shape(x, 3) // inC

というわけで この strides attributes の値も [0, 2, 3, 1] の順に並び換えたものを用意する。

ここまで準備できたらいよいよ operation の置き換え。

    # New operations
    new_op = graph.create_op(op_type=target_op.type, inputs=inputs, name=f'{target_op.name}_nhwc', attrs=attrs)
    output = tf.transpose(new_op.outputs[0], [0, 3, 1, 2], name=f'{new_op.name}_output')

    # Update connections
    ops = [op for op in graph.get_operations() if target_op.outputs[0] in op.inputs]
    for op in ops:
        for i, input_tensor in enumerate(op.inputs):
            if input_tensor.name == target_op.outputs[0].name:
                op._update_input(i, output)

元の operation と同じ type, attrs を持つ新しい operation を作成する。入力は NHWC に transpose したもの。 ということは この operation の outputs も NHWC になっているので、その後の計算に支障が出ないよう この出力は NCHW に戻しておいてやる必要がある。ので outputs を今度は [0, 3, 1, 2] でtransposeする。

そして、元々の outputs を受け取っていた 次の operation たちの入力を この新しい operation の outputs に置き換えてやる。

PrevOp ---> NCHW ---------> Op(NCHW) -----------NCHW-----> NextOp
        |                                              |
         -> NCHW to NHWC -> Op(NHWC) -> NHWC to NCHW --

元々上段の流れだけだったものに対して、下段のルートを付け足した形になる。

最後に、元々あった NCHW の operation を graph から消しておく。

# Delete old nodes
graph_def = graph.as_graph_def()
for target_op in target_ops:
    graph_def.node.remove(target_op.node_def)

これで、一度 SavedModel に graph を保存してみよう。

inputs = graph.get_tensor_by_name(f'Gs/{Gs.input_names[0]}:0')
outputs = graph.get_tensor_by_name(f'Gs/{Gs.output_names[0]}:0')

tf.compat.v1.enable_resource_variables()
tf.compat.v1.saved_model.simple_save(
    tf.get_default_session(),
    './savedmodel',
    {'inputs': inputs},
    {'outputs': outputs},
)

これを load して実行してみると…

import tensorflow as tf

model = tf.compat.v2.saved_model.load('./savedmodel')
generate = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

z = tf.random.normal([1, 512])
outputs = generate(inputs=z)['outputs']
with tf.compat.v1.Session() as sess:
    sess.run(tf.compat.v1.initializers.tables_initializer())
    print(sess.run(outputs))
2020-02-04 23:24:38.004217: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2020-02-04 23:24:38.014420: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x7fbdbabd65a0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
2020-02-04 23:24:38.014438: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Host, Default Version
[[[[0.65079993 0.7005807  0.71635807 ... 0.70904994 0.672469
    0.65049875]
   [0.7084702  0.73016286 0.7354313  ... 0.74900943 0.7347464
    0.7381906 ]
   [0.71451813 0.7372785  0.73747885 ... 0.75236607 0.7619566
    0.7417464 ]
   ...

何らかの数値が出力された!

これはやはり NCHW(1, 3, 256, 256) のような形で来ているので、RGBの画像として Pillow などで扱うにはやはり NHWC に transpose したりといった処理は必要になる。

import tensorflow as tf
from PIL import Image

model = tf.compat.v2.saved_model.load('./savedmodel')
generate = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

z = tf.random.normal([1, 512])
outputs = generate(inputs=z)['outputs']
outputs = tf.transpose(outputs, [0, 2, 3, 1])
outputs = tf.saturate_cast((outputs + 1.0) * 127.5, tf.uint8)
with tf.compat.v1.Session() as sess:
    sess.run(tf.compat.v1.initializers.tables_initializer())
    img = Image.fromarray(sess.run(outputs)[0])
img.save('output.png')

f:id:sugyan:20200204233514p:plain

CPU環境でも生成が出来た! やったー!!

2. Modelを構築し直す

標準的な畳み込みを使ったネットワークであれば ここまでの変換で問題なく動くようになるかもしれない。が、StyleGAN2の場合はまだちょっとだけ問題が残っている。

単一の画像生成なら上述のように出来るけど、 batch_size1 より大きくすると、またエラーになってしまう。

z = tf.random.normal([2, 512])
...

tensorflow.python.framework.errors_impl.UnimplementedError: [_Derived_]{{function_node __inference_pruned_14657}} {{function_node __inference_pruned_14657}} The Conv2D op currently does not support grouped convolutions on the CPU. A grouped convolution was attempted to be run because the input depth of 1024 does not match the filter input depth of 512
         [[{{node Gs/G_synthesis/4x4/Conv/Conv2D_nhwc}}]]
         [[StatefulPartitionedCall_1]]

grouped convolution というものが使われていて、これがまた CPU環境ではまだサポートされていないものだった。

このへんかな?

TensorFlow 1.14 あたりから入ったもののようだ。 普通は filter のshapeは [filter_height, filter_width, in_channels, out_channels] と なっていて、inputchannelfilter2 番目の次元と等しくなければならないが、cuDNNによって filter.shape[2] の倍数であればまとめて計算できるようになる、という機能… なのかな。(よく分かっていない)

これがどうやら networks_stylegan2.modulated_conv2d_layer() の中で fused_modconv=True のときにそういった処理をするようになっているらしい。

networks_stylegan2.py 抜粋:

    # Reshape/scale input.
    if fused_modconv:
        x = tf.reshape(x, [1, -1, x.shape[2], x.shape[3]]) # Fused => reshape minibatch to convolution groups.
        w = tf.reshape(tf.transpose(ww, [1, 2, 3, 0, 4]), [ww.shape[1], ww.shape[2], ww.shape[3], -1])
    else:
        x *= tf.cast(s[:, :, np.newaxis, np.newaxis], x.dtype) # [BIhw] Not fused => scale input activations.

    # Convolution with optional up/downsampling.
    if up:
        x = upsample_conv_2d(x, tf.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
    elif down:
        x = conv_downsample_2d(x, tf.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
    else:
        x = tf.nn.conv2d(x, tf.cast(w, x.dtype), data_format='NCHW', strides=[1,1,1,1], padding='SAME')

    # Reshape/scale output.
    if fused_modconv:
        x = tf.reshape(x, [-1, fmaps, x.shape[2], x.shape[3]]) # Fused => reshape convolution groups back to minibatch.

up/downsampling の処理をかける前に xw をreshapeして、その結果を後でまたreshapeして元に戻す、といった形になっているようだ。

これは単一の operation の書き換えでは対応できない…。


ということで、結局 graph の書き換えだけでは無理そうなので 「重みだけ再利用してModelはCPUでも動かせる形に構築し直す」という方針を取ることにした。

まずは .pkl をloadした後、 checkpoint 形式で 変数の値だけを保存する。

import pickle
import tensorflow as tf
from dnnlib import tflib

tflib.init_tf()
with open('network.pkl', 'rb') as fp:
    _, _, Gs = pickle.load(fp, encoding='latin1')

saver = tf.compat.v1.train.Saver(Gs.vars)
saver.save(tf.get_default_session(), './ckpt/network')
$ ls -l ./ckpt
total 240176
-rw-r--r--  1 sugyan  staff         71 Feb  4 23:45 checkpoint
-rw-r--r--  1 sugyan  staff  120838348 Feb  4 23:45 network.data-00000-of-00001
-rw-r--r--  1 sugyan  staff       4898 Feb  4 23:45 network.index
-rw-r--r--  1 sugyan  staff    2114457 Feb  4 23:45 network.meta

で、 fused_modconv=False になるような設定で Generator を作成する。 どうせ graph を構築しなおすことになるのだし、前述したような NCHW -> NHWC の書き換えもコード上でやってしまおう。

Gs/ に関係するところだけなら 以下の3箇所の tf.nn.conv2dtf.nn.conv2d_transpose の入出力まわりを書き換えてやれば大丈夫だ。

diff --git a/dnnlib/tflib/ops/upfirdn_2d.py b/dnnlib/tflib/ops/upfirdn_2d.py
index fd23777..26cc573 100755
--- a/dnnlib/tflib/ops/upfirdn_2d.py
+++ b/dnnlib/tflib/ops/upfirdn_2d.py
@@ -93,7 +93,9 @@ def _upfirdn_2d_ref(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):
     x = tf.transpose(x, [0, 3, 1, 2])
     x = tf.reshape(x, [-1, 1, inH * upy + pady0 + pady1, inW * upx + padx0 + padx1])
     w = tf.constant(k[::-1, ::-1, np.newaxis, np.newaxis], dtype=x.dtype)
-    x = tf.nn.conv2d(x, w, strides=[1,1,1,1], padding='VALID', data_format='NCHW')
+    x = tf.transpose(x, [0, 2, 3, 1])
+    x = tf.nn.conv2d(x, w, strides=[1,1,1,1], padding='VALID', data_format='NHWC')
+    x = tf.transpose(x, [0, 3, 1, 2])
     x = tf.reshape(x, [-1, minorDim, inH * upy + pady0 + pady1 - kernelH + 1, inW * upx + padx0 + padx1 - kernelW + 1])
     x = tf.transpose(x, [0, 2, 3, 1])

@@ -288,7 +290,11 @@ def upsample_conv_2d(x, w, k=None, factor=2, gain=1, data_format='NCHW', impl='c
     w = tf.reshape(w, [convH, convW, -1, num_groups * inC])

     # Execute.
-    x = tf.nn.conv2d_transpose(x, w, output_shape=output_shape, strides=stride, padding='VALID', data_format=data_format)
+    x = tf.transpose(x, [0, 2, 3, 1])
+    stride = [1, factor, factor, 1]
+    output_shape = [_shape(x, 0), (_shape(x, 1) - 1) * factor + convH, (_shape(x, 2) - 1) * factor + convW, outC]
+    x = tf.nn.conv2d_transpose(x, w, output_shape=output_shape, strides=stride, padding='VALID', data_format='NHWC')
+    x = tf.transpose(x, [0, 3, 1, 2])
     return _simple_upfirdn_2d(x, k, pad0=(p+1)//2+factor-1, pad1=p//2+1, data_format=data_format, impl=impl)

 #----------------------------------------------------------------------------
diff --git a/training/networks_stylegan2.py b/training/networks_stylegan2.py
index 6c96fc1..8fe2979 100755
--- a/training/networks_stylegan2.py
+++ b/training/networks_stylegan2.py
@@ -117,7 +117,9 @@ def modulated_conv2d_layer(x, y, fmaps, kernel, up=False, down=False, demodulate
     elif down:
         x = conv_downsample_2d(x, tf.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
     else:
-        x = tf.nn.conv2d(x, tf.cast(w, x.dtype), data_format='NCHW', strides=[1,1,1,1], padding='SAME')
+        x = tf.transpose(x, [0, 2, 3, 1])
+        x = tf.nn.conv2d(x, tf.cast(w, x.dtype), data_format='NHWC', strides=[1,1,1,1], padding='SAME')
+        x = tf.transpose(x, [0, 3, 1, 2])

     # Reshape/scale output.
     if fused_modconv:

で、 Generator を作成し、保存しておいた変数を checkpoint ファイルからloadする。

import tensorflow as tf
from dnnlib import tflib
from dnnlib import EasyDict

tflib.init_tf()
G_args = EasyDict(func_name='training.networks_stylegan2.G_main')
G_args.fused_modconv = False
G = tflib.Network(
    'G',
    num_channels=3,
    resolution=256,
    **G_args)
Gs = G.clone('Gs')

saver = tf.compat.v1.train.Saver(Gs.vars)
saver.restore(tf.get_default_session(), 'ckpt/network')

graph = tf.get_default_graph()
inputs = graph.get_tensor_by_name(f'Gs/{Gs.input_names[0]}:0')
outputs = graph.get_tensor_by_name(f'Gs/{Gs.output_names[0]}:0')
tf.compat.v1.enable_resource_variables()
tf.compat.v1.saved_model.simple_save(
    tf.get_default_session(),
    './savedmodel',
    {'inputs': inputs},
    {'outputs': outputs}
)

これで今度は NCHW の operation も fused_modconv による grouped convolution も含まない SavedModel が出力された、はず。

再度 batch_size > 1 で生成をしてみよう。

model = tf.compat.v2.saved_model.load('./savedmodel')
generate = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]

z = tf.random.normal([3, 512])
outputs = generate(inputs=z)['outputs']
outputs = tf.transpose(outputs, [0, 2, 3, 1])
outputs = tf.saturate_cast((outputs + 1.0) * 127.5, tf.uint8)
with tf.compat.v1.Session() as sess:
    sess.run(tf.compat.v1.initializers.tables_initializer())
    img = Image.fromarray(np.concatenate(sess.run(outputs), axis=1))
img.save('output.png')

f:id:sugyan:20200205232314p:plain

ちゃんと3つ画像が生成された! やったー!!

TensorFlow.jsで動かす

さて、ここまで出来ているのならば特殊な operation なども使っていないはずだろうし TensorFlow.js の GraphModel に変換できるだろう、と思うわけです。

$ tensorflowjs_converter \
    --input_format=tf_saved_model \
    --output_format=tfjs_graph_model \
    --signature_name=serving_default \
    --saved_model_tags=serve \
    --skip_op_check \
    ./savedmodel ./tfjs

やはり RandomStandardNormal の operation はサポートされていないので これがある場合は --skip_op_check optionが必要になる。

ちなみに random入力を使っている場所は randomize_noise というパラメータで制御されていて、 Generator の作成時にこれを fused_modconv と同様に G_args.randomize_noise = False と指定しておくとこの operation は使われなくなって、 --skip_op_check でskipする必要が無くなる。

ransomize_noise がどれくらい出力の品質に影響あるかちょっとよく分かっていないけど、問題ないようなら False にしておけば計算の負荷も減るし良いかもしれない。

ともかく、変換した GraphModel を読み込んで実行してみると…

const url = '.../tfjs/model.json'
const randomNormal = (node) => {
    return tf.randomNormal(node.inputs[0].dataSync())
};
tf.registerOp('RandomStandardNormal', randomNormal);
tf.loadGraphModel(url).then((model) => {
    tf.tidy(() => {
        const z = tf.randomNormal([1, 512])
        const results = model.execute(z)
        console.log(results.shape)
    })
}).catch((err) => {
    console.error(err);
})
webgl_util.ts:110 Uncaught Error: Failed to compile fragment shader.

やはり webgl backend ではエラーが出てしまう。 試しに tf.setBackend('cpu') にしてみると、ものすごい時間はかかるが 一応実行は可能なようだ。しかし現実的ではない。

何故 webgl backend では上手くいかないのか… とひたすら地道に探っていたところ、一つ 特殊な箇所を発見した。

upfirdn_2d の reference implementation で、Upsampleするために Rank 6Tensorreshape してから pad を行っている場所がある。

upfirdn_2d.py 抜粋:

def _upfirdn_2d_ref(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):
    """Slow reference implementation of `upfirdn_2d()` using standard TensorFlow ops."""

    ...

    # Upsample (insert zeros).
    x = tf.reshape(x, [-1, inH, 1, inW, 1, minorDim])
    x = tf.pad(x, [[0, 0], [0, 0], [0, upy - 1], [0, 0], [0, upx - 1], [0, 0]])
    x = tf.reshape(x, [-1, inH * upy, inW * upx, minorDim])

    ...

この関数の処理の詳細は理解できていないけど、ともかくこの箇所では 0 で padding することにより Tensor のサイズを大きくしていることだけは分かる。

ところで TensorFlow.js の tf.padAPI reference を見てみると…

Also available are stricter rank-specific methods with the same signature as this method that assert that paddings is of given length.

  • tf.pad1d
  • tf.pad2d
  • tf.pad3d
  • tf.pad4d

…これ、Rank 5 以上のものには対応していないのでは!?

軽く tensorflow/tfjs のコードを見てみたがちょっと分からず… しかしまぁ Rank 4 までしか対応していない、というのは実に有り得る気がする。

ということで 前述の Rank 6Tensor に対する pad の回避するよう処理を書き換えてみることにした。

とはいえ data_format のときのように transpose すれば良いというものでもなさそうだし ちょっとどうすれば良いか分からない…。

が、幸い ここで Upsample するためのパラメータ upy, upx はどうやら 1 もしくは 2 の値でしか渡されてこないらしい、ということが分かった。 1 のときは結局 pad すべき shape は 0 になるので、何もせずに skip してしまえば良い。 2 のときだけ どうにか2箇所だけ shape を増やしてあげる必要がある…。

いや、待てよこれは 1列分だけ 0 padding するだけなのだから、同じ shape の zeros Tensor を後ろから concat してやれば同じ意味になるのでは? と思いついた。

diff --git a/dnnlib/tflib/ops/upfirdn_2d.py b/dnnlib/tflib/ops/upfirdn_2d.py
index fd23777..49d11ed 100755
--- a/dnnlib/tflib/ops/upfirdn_2d.py
+++ b/dnnlib/tflib/ops/upfirdn_2d.py
@@ -82,7 +82,10 @@ def _upfirdn_2d_ref(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):

     # Upsample (insert zeros).
     x = tf.reshape(x, [-1, inH, 1, inW, 1, minorDim])
-    x = tf.pad(x, [[0, 0], [0, 0], [0, upy - 1], [0, 0], [0, upx - 1], [0, 0]])
+    if upy == 2:
+        x = tf.concat([x, tf.zeros(tf.shape(x))], axis=2)
+    if upx == 2:
+        x = tf.concat([x, tf.zeros(tf.shape(x))], axis=4)
     x = tf.reshape(x, [-1, inH * upy, inW * upx, minorDim])

     # Pad (crop if negative).

このように書き換える。

これでまた Generator を作成しなおして SavedModel に保存して GraphModel に変換して…。

f:id:sugyan:20200206005109p:plain

TensorFlow.js でも 動いた!! やったー!!

やはり webgl backend でのエラーは tf.pad が Rank 4 までしか対応していなかったことが原因だったようだ…。いやーまさかこんな方法で解決できるとは。。

実際のところ、 TensorFlow.js での生成を 256x256 サイズで試した感じでは 最初の実行時に 10秒弱かかるが、その後は 数十 ms くらいで計算できてそう。そこから計算結果のデータを取得して、という部分で 3000 ms くらいかかってしまっているが…。

とりあえず 計算はWebWorkerに任せる などして、描画とかUIのところだけ作っていけば、 誰でもブラウザ上で画像生成を試せるようになる、かも…!? という希望は見えた。

ここまでの変更は一応 ここに残しておく。