fidjax

FID ジャックス

JAX での Frechet Inception Distance のクリーンな実装。

• OpenAI の TensorFlow 実装を再現します。
• 純粋な JAX 実装は、CPU/GPU/TPU および JIT 内で実行されます。
• pathlib API を使用して GCS から重みをロードできます。
• クリーンでシンプルなコード。

説明書

1️⃣ FID JAX は単一のファイルなので、プロジェクト ディレクトリにコピーするだけです。または、パッケージをインストールすることもできます。

pip install fidjax

2️⃣ インセプション ウェイトをダウンロードします (Matthias Wright のクレジット):

wget https://www.dropbox.com/s/xt6zvlvt22dcwck/inception_v3_weights_fid.pickle ? dl=1

3️⃣ 希望の解像度の ImageNet 参照統計をダウンロードします (他のデータセット用に独自の統計を生成します)。

wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/64/VIRTUAL_imagenet64_labeled.npz
wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/128/VIRTUAL_imagenet128_labeled.npz
wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/VIRTUAL_imagenet256_labeled.npz
wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/512/VIRTUAL_imagenet512.npz

4️⃣ JAX でアクティベーション、統計、スコアを計算します。

 import fidjax
import numpy as np

weights = './inception_v3_weights_fid.pickle?dl=1'
reference = './VIRTUAL_imagenet128_labeled.npz'
fid = fidjax . FID ( weights , reference )

fid_total = 50000
fid_batch = 1000
acts = []
for range ( fid_total // fid_batch ):
  samples = ...  # (B, H, W, 3) jnp.uint8
  acts . append ( fid . compute_acts ( samples ))
stats = fid . compute_stats ( acts )
score = fid . compute_score ( stats )

print ( float ( score ))  # FID 

正確さ

データセット	モデル	FID ジャックス	OpenAI TF
イメージネット 256	ADM (ガイド付き、アップサンプリング)	3.937	3.943

チュートリアル

クラウドストレージの使用

クラウド ストレージをサポートするpathlib.Path実装を介してファイルを指します。たとえば GCS の場合:

 import elements  # pip install elements
import fidjax

weights = elements . Path ( 'gs://bucket/fid/inception_v3_weights_fid.pickle' )
reference = elements . Path ( 'gs://bucket/fid/VIRTUAL_imagenet128_labeled.npz' )

fid = fidjax . FID ( weights , reference )

カスタム データセット

カスタム データセットの参照統計を生成します。

 import fidjax
import numpy as np

weights = './inception_v3_weights_fid.pickle?dl=1'
fid = fidjax . FID ( weights )

acts = fid . compute_acts ( images )
mu , sigma = fid . compute_stats ( acts )

np . savez ( 'reference.npz' , { 'mu' : mu , 'sigma' : sigma })

リソース

• TensorFlow での OpenAI 標準実装
• PyTorch ポート
• JAXポート
• JAX ポート (単一ファイル)

質問

Github に問題を提出してください。