fidjax

FID賈克斯

JAX 中 Frechet 起始距離的乾淨實現。

• 重現 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 上提交問題。