TensorFlow 2.0 の Frechet Inception Distance (FID) の計算の tf.linalg.sqrtm() で Nan が出るときの応急処置

概要

[1] の numpy で書かれたコードを TensorFlow 2.0 に移植するため、FID の計算を TensorFlow のみでしようとしていたところ、scipy.linalg.sqrtm() では計算可能な 2048x2048 の ndarray の平方根を求めようとしても、TensorFlow 側だけ Nan になってしまう現象にぶつかった。

• TensorFlow 1.0 でも同じ現象が起こるというイシューが GitHub に立っている (同じことをすれば対処できるはず)

対処法

対処には二段階必要です。

まず、tf.linalg.sqrtm() の一個前の計算である、$\Sigma_r \Sigma_g$ の時点で、tf.cast を噛ませて float64 に設定しないといけません。TensorFlow のデフォルトは float32 のため、計算が失敗する可能性があるらしい¹です。

sigma1, sigma2 = tf.cast(sigma1, tf.float64), tf.cast(sigma2, tf.float64)

次に、[2] を参考に、微小の数値 eps を対角行列として加え、その後行列の掛け算を行います。ここでは、全て型を tf.float64 に設定しないといけません。

eps = tf.constant(1e-6, dtype=tf.float64)
offset = tf.eye(2048, dtype=tf.float64) * eps
tdot = tf.tensordot(sigma1+offset, sigma2+offset, axes=1)

これで対処が終わりました。あとは、

covmean = tf.linalg.sqrtm(tdot)

で、一応計算が通り、ほぼ同じ計算結果になります。

問題点

計算は止まらなくなったものの、TensorFlow 側の計算が、Scipy の計算より数十倍遅いので、修正方法が分かり次第追記する。

参考

• [1] How to Implement the Frechet Inception Distance (FID) for Evaluating GANs: https://machinelearningmastery.com/how-to-implement-the-frechet-inception-distance-fid-from-scratch/
• [2] FID get "nan" or "complex number": https://github.com/bioinf-jku/TTUR/issues/4

1. https://github.com/tensorflow/tensorflow/issues/31731#issuecomment-522617924 ↩