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のところだけ作っていけば、 誰でもブラウザ上で画像生成を試せるようになる、かも…!? という希望は見えた。

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

StyleGAN2による画像生成をJSで動かしたい

memo.sugyan.com

ということで 多少集めることができた画像データを使って、StyleGAN2を使って生成を試してみた。

github.com

こんな感じで学習が進み、ある程度はそれっぽい顔画像が生成できるようになってきた。

f:id:sugyan:20200129231329p:plain

これはまだこれから学習データもさらに増やして改善していくとして、それはそれとして こうした生成をWebブラウザ上でも動かせるようにしたい。 と思った。

とにかくやってみないと分からん、ってことで挑戦。

Pickle to SavedModel

まず、学習の過程で保存されるsnapshot は network-snapshot-000080.pkl のような形で 学習中のネットワーク全体を pickle で固めたものとして保存されている。

これを持ってきて解凍してみるところから。 run_generator.py をみると、以下のように読み込んでいる。

from dnnlib import tflib

tflib.init_tf()
network_pkl = 'network-snapshot-000080.pkl'
with open(network_pkl, 'rb') as fp:
    _G, _D, Gs = pickle.load(fp, encoding='latin1')

しかしこれはGPUが使える環境じゃないと この読み込みさえエラーになってしまう。 dnnlib/tflib/network.py__setstate__ でgraphの構築をしていて、その中で fused_bias_act, upfirdn_2d といったcustom opsを使用している。 これらが、デフォルトでCUDAのAPIを使ったimplementationを使用しているから、のようだ。

もしCPU環境でmodelを読み込もうとしたら、これらをどうにかする必要がある。幸い CUDAを使わない実装も用意されているので

-def fused_bias_act(x, b=None, axis=1, act='linear', alpha=None, gain=None, impl='cuda'):
+def fused_bias_act(x, b=None, axis=1, act='linear', alpha=None, gain=None, impl='ref'):

のように引数のdefault impl'ref' に変えてやることで CUDAが使えない状態でも問題なくloadは出来るようになる。

で、ともかくloadできた Gs から inputs, outputs を取得して SavedModel の形式で保存するようにしてみた。

import pickle
import tensorflow as tf
from dnnlib import tflib


builder = tf.saved_model.Builder('savedmodel')
with tf.Graph().as_default() as graph:
    tflib.init_tf()
    network_pkl = 'network-snapshot-000080.pkl'
    with open(network_pkl, 'rb') as fp:
        _, _, Gs = pickle.load(fp, encoding='latin1')

    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')
    signature = tf.saved_model.build_signature_def(
        {'inputs': tf.compat.v1.saved_model.build_tensor_info(inputs)},
        {'outputs': tf.compat.v1.saved_model.build_tensor_info(outputs)})
    builder.add_meta_graph_and_variables(
        tf.get_default_session(),
        [tf.saved_model.SERVING],
        signature_def_map={
            tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature
        })
builder.save()

これで、pklのsnapshotからSavedModelの形式に変換することが出来た、ということになる。

$ saved_model_cli show --all --dir savedmodel

MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['inputs'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 512)
        name: Gs/latents_in:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['outputs'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 3, 256, 256)
        name: Gs/images_out:0
  Method name is:

試しにこのSavedModelを読み込んで生成を実行してみる。

import tensorflow as tf
import numpy as np

with tf.Graph().as_default() as graph:
    model = tf.compat.v2.saved_model.load('./savedmodel')

    generate = model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
    rnd = np.random.RandomState(0)
    z = tf.constant(rnd.randn(1, 512), tf.float32)

    with tf.compat.v1.Session(graph=graph) as sess:
        sess.run(tf.compat.v1.initializers.global_variables())
        print(sess.run(generate(inputs=z)['outputs']))

すると

2020-01-29 23:33:16.173370: 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}}]]
Traceback (most recent call last):

...

NHWC でしか動かないものがあって でもmodelは NCHW で作られて固められてしまっているのでダメなようだ…

SavedModel to GraphModel

SavedModelそのままではCPU環境では動かせないのだけど、とりあえず気にせずJSのmodelとしてconvertしてみる。もしかしたらoptimizeの過程でそのへんもどうにかしてくれるかも? と期待してみたり。

pip install tensorflowjs して convertをかけてみる。

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

山ほど警告が出た後に、エラーで止まる。

...

ValueError: Unsupported Ops in the model before optimization
RandomStandardNormal

どうやら tensorflowjs_converter では RandomStandardNormal などのopsはサポートされていないらしい。困った。 しかし調べてみると --skip_op_check というoptionが用意されていて、サポートしてないopがあってもとりあえずskipしてくれるようだ。

$ 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

...

2020-01-29 23:43:17.764627: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:786] Optimization results for grappler item: graph_to_optimize
2020-01-29 23:43:17.764712: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   debug_stripper: debug_stripper did nothing. time = 0.969ms.
2020-01-29 23:43:17.764717: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   model_pruner: Graph size after: 2022 nodes (-115), 2099 edges (-115), time = 187.757ms.
2020-01-29 23:43:17.764721: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   constant_folding: Graph size after: 1614 nodes (-408), 1691 edges (-408), time = 1340.328ms.
2020-01-29 23:43:17.764725: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   arithmetic_optimizer: Graph size after: 962 nodes (-652), 1655 edges (-36), time = 209.671ms.
2020-01-29 23:43:17.764728: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   dependency_optimizer: Graph size after: 854 nodes (-108), 1453 edges (-202), time = 108.008ms.
2020-01-29 23:43:17.764805: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   model_pruner: Graph size after: 854 nodes (0), 1453 edges (0), time = 141.52ms.
2020-01-29 23:43:17.764815: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   constant_folding: Graph size after: 854 nodes (0), 1453 edges (0), time = 369.232ms.
2020-01-29 23:43:17.764819: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   arithmetic_optimizer: Graph size after: 854 nodes (0), 1453 edges (0), time = 197.42ms.
2020-01-29 23:43:17.764823: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   dependency_optimizer: Graph size after: 842 nodes (-12), 1441 edges (-12), time = 76.729ms.
2020-01-29 23:43:17.764826: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   debug_stripper: debug_stripper did nothing. time = 10.353ms.
2020-01-29 23:43:17.764829: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   model_pruner: Graph size after: 842 nodes (0), 1441 edges (0), time = 123.947ms.
2020-01-29 23:43:17.764993: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   constant_folding: Graph size after: 842 nodes (0), 1441 edges (0), time = 365.53ms.
2020-01-29 23:43:17.765002: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   arithmetic_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 191.759ms.
2020-01-29 23:43:17.765006: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   dependency_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 77.682ms.
2020-01-29 23:43:17.765010: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   model_pruner: Graph size after: 842 nodes (0), 1441 edges (0), time = 129.848ms.
2020-01-29 23:43:17.765013: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   constant_folding: Graph size after: 842 nodes (0), 1441 edges (0), time = 360.265ms.
2020-01-29 23:43:17.765016: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   arithmetic_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 212.4ms.
2020-01-29 23:43:17.765073: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   dependency_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 79.051ms.
2020-01-29 23:43:22.461213: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:786] Optimization results for grappler item: graph_to_optimize
2020-01-29 23:43:22.461239: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   remapper: Graph size after: 842 nodes (0), 1441 edges (0), time = 149.812ms.
2020-01-29 23:43:22.461244: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   constant_folding: Graph size after: 842 nodes (0), 1441 edges (0), time = 370.785ms.
2020-01-29 23:43:22.461248: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   arithmetic_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 187.354ms.
2020-01-29 23:43:22.461251: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   dependency_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 78.03ms.
2020-01-29 23:43:22.461255: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   remapper: Graph size after: 842 nodes (0), 1441 edges (0), time = 136.545ms.
2020-01-29 23:43:22.461258: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   constant_folding: Graph size after: 842 nodes (0), 1441 edges (0), time = 370.468ms.
2020-01-29 23:43:22.461343: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   arithmetic_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 199.41ms.
2020-01-29 23:43:22.461353: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:788]   dependency_optimizer: Graph size after: 842 nodes (0), 1441 edges (0), time = 79.184ms.
Writing weight file ./tfjs/model.json...

$ ls ./tfjs
group1-shard10of29.bin
group1-shard11of29.bin
group1-shard12of29.bin
group1-shard13of29.bin
group1-shard14of29.bin
group1-shard15of29.bin
group1-shard16of29.bin
group1-shard17of29.bin
group1-shard18of29.bin
group1-shard19of29.bin
group1-shard1of29.bin
group1-shard20of29.bin
group1-shard21of29.bin
group1-shard22of29.bin
group1-shard23of29.bin
group1-shard24of29.bin
group1-shard25of29.bin
group1-shard26of29.bin
group1-shard27of29.bin
group1-shard28of29.bin
group1-shard29of29.bin
group1-shard2of29.bin
group1-shard3of29.bin
group1-shard4of29.bin
group1-shard5of29.bin
group1-shard6of29.bin
group1-shard7of29.bin
group1-shard8of29.bin
group1-shard9of29.bin
model.json

何かが出力された。めでたい。

Load GraphModel and Execute

無事(?)にTensorFlow.jsの GraphModel形式にconvert出来たので、ブラウザ上でこれをloadして実行してみる。

const url = 'http://localhost:8000/tfjs/model.json'
tf.loadGraphModel(url).then((model) => {
    const z = tf.randomUniform([1, 512]);
    model.execute(z).print();
}).catch((err) => {
    console.error(err);
})

するとConsoleで以下のようなエラーが。

 TypeError: Unknown op 'RandomStandardNormal'. File an issue at https://github.com/tensorflow/tfjs/issues so we can add it, or register a custom execution with tf.registerOp()

はいはい。確かにサポートしてないopsだからskipしてもらってたヤツだ。

しかしどうやら tf.registerOp() で自分で定義すれば良いらしい。 tf.randomNormal というAPIはあるので、それがそのまま使えるはず。

const randomNormal = (node) => {
    return tf.randomNormal(node.inputs[0].shape);
};
tf.registerOp('RandomStandardNormal', randomNormal);

というのを足してみる。

そうするとそれっぽく実行できる雰囲気になったのだけど、今度は webgl_util.ts 内で

Error: Failed to compile fragment shader.

という現象が起きて失敗…。えー、、なにそれ。。。

webgl backendでダメなのかな、、と tf.setBackend('cpu') してみると、今度は

Error: Error in conv2dDerInput: depth of input (9) must match input depth for filter 512.

と別のエラー。shapeが合わない…? NCHW のものによる弊害がここで出てきてるのかな…?

現状

ここまでで止まってしまっている状態。JSのmodelをloadするところまでは出来ているから もう一息かな〜と思ったのだけど。。。

modelを一度 NHWC の形式に変換してやるとかの処理が必要になるんだろうか… うーん