EDM2
1.0.0
拡散モデルのトレーニングダイナミクスの分析と改善
Tero Karras、Miika Aittala、Jaakko Lehtinen、Janne Hellsten、Timo Aila、Samuli Laine
https://arxiv.org/abs/2312.02696
重要な貢献：
conda env create -f environment.yml -n edmconda activate edm私たちの論文の主な結果を再現するには、単に実行してください。
python example.pyこれは、各データセットに最適な事前トレーニングモデルをロードし、最適なサンプラー設定を使用して画像のランダムな8x8グリッドを生成する最小限のスタンドアロンスクリプトです。期待される結果：
| データセット | ランタイム | 参照画像 |
|---|---|---|
| CIFAR-10 | 〜6秒 | cifar10-32x32.png |
| ffhq | 〜28秒 | ffhq-64x64.png |
| AFHQV2 | 〜28秒 | afhqv2-64x64.png |
| Imagenet | 〜5分 | imagenet-64x64.png |
さまざまなサンプリング戦略を探索する最も簡単な方法は、 example.py直接変更することです。また、関連するビットをコピーするだけで、事前に訓練されたモデルおよび/または提案されたEDMサンプラーを独自のコードに組み込むこともできます。事前に訓練されたモデルのクラス定義は、ピクルス自体に保存され、 torch_utils.persistenceを介して解読中に自動的にロードされることに注意してください。外部Pythonスクリプトでモデルを使用するには、 PYTHONPATHを通じてtorch_utilsとdnnlibアクセスできることを確認してください。
Docker ：次のようにDockerを使用してサンプルスクリプトを実行できます。
# Build the edm:latest image
docker build --tag edm:latest .
# Run the generate.py script using Docker:
docker run --gpus all -it --rm --user $( id -u ) : $( id -g )
-v ` pwd ` :/scratch --workdir /scratch -e HOME=/scratch
edm:latest
python example.py注：Dockerイメージには、NVIDIAドライバーがr520以降をリリースする必要があります。
docker run Inkocationは気が遠くなるように見えるかもしれないので、ここにその内容を解きましょう。
--gpus all -it --rm --user $(id -u):$(id -g) ：すべてのGPUを有効にして、現在のユーザーのUID/GIDとのインタラクティブセッションを実行して、Dockerがルートとしてファイルを作成することを避けます。-v `pwd`:/scratch --workdir /scratch ：マウントカレントランニングdir（例：ホストマシンのこのgitリポジトリの上部）をコンテナ内で/scratch 、それを現在の作業担当者として使用します。-e HOME=/scratch ：一時ファイルをキャッシュする場所を指定します。注：より微細なコントロールが必要な場合は、代わりにDNNLIB_CACHE_DIR設定できます（事前に訓練されたモデルのダウンロードキャッシュの場合）。これらのキャッシュ監督は、コンテンツが複数のdocker run呼び出しで保持されるように、永続的なボリュームに存在するようにします。 提案されたトレーニング構成（構成F）およびベースライン構成（構成A）の事前に訓練されたモデルを提供します。
特定のモデルとサンプラーを使用して画像のバッチを生成するには、実行します。
# Generate 64 images and save them as out/*.png
python generate.py --outdir=out --seeds=0-63 --batch=64
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-cond-vp.pkl多数の画像を生成するのは時間がかかる場合があります。 torchrunを使用して上記のコマンドを起動することにより、ワークロードを複数のGPUに分配できます。
# Generate 1024 images using 2 GPUs
torchrun --standalone --nproc_per_node=2 generate.py --outdir=out --seeds=0-999 --batch=64
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-cond-vp.pklサンプラー設定は、コマンドラインオプションを使用して制御できます。詳細については、 python generate.py --help参照してください。最良の結果を得るには、各データセットに次の設定を使用することをお勧めします。
# For CIFAR-10 at 32x32, use deterministic sampling with 18 steps (NFE = 35)
python generate.py --outdir=out --steps=18
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-cond-vp.pkl
# For FFHQ and AFHQv2 at 64x64, use deterministic sampling with 40 steps (NFE = 79)
python generate.py --outdir=out --steps=40
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-ffhq-64x64-uncond-vp.pkl
# For ImageNet at 64x64, use stochastic sampling with 256 steps (NFE = 511)
python generate.py --outdir=out --steps=256 --S_churn=40 --S_min=0.05 --S_max=50 --S_noise=1.003
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-imagenet-64x64-cond-adm.pkl提案されているEDMサンプラーに加えて、 generate.py使用して、ペーパーのセクション3のサンプラーアブレーションを再現することもできます。例えば：
# Figure 2a, "Our reimplementation"
python generate.py --outdir=out --steps=512 --solver=euler --disc=vp --schedule=vp --scaling=vp
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/baseline/baseline-cifar10-32x32-uncond-vp.pkl
# Figure 2a, "+ Heun & our {t_i}"
python generate.py --outdir=out --steps=128 --solver=heun --disc=edm --schedule=vp --scaling=vp
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/baseline/baseline-cifar10-32x32-uncond-vp.pkl
# Figure 2a, "+ Our sigma(t) & s(t)"
python generate.py --outdir=out --steps=18 --solver=heun --disc=edm --schedule=linear --scaling=none
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/baseline/baseline-cifar10-32x32-uncond-vp.pkl特定のモデルとサンプラーのFréchetInception距離（FID）を計算するには、最初に50,000のランダム画像を生成し、次にfid.pyを使用してデータセット参照統計と比較してください。
# Generate 50000 images and save them as fid-tmp/*/*.png
torchrun --standalone --nproc_per_node=1 generate.py --outdir=fid-tmp --seeds=0-49999 --subdirs
--network=https://nvlabs-fi-cdn.nvidia.com/edm/pretrained/edm-cifar10-32x32-cond-vp.pkl
# Calculate FID
torchrun --standalone --nproc_per_node=1 fid.py calc --images=fid-tmp
--ref=https://nvlabs-fi-cdn.nvidia.com/edm/fid-refs/cifar10-32x32.npz上記の両方のコマンドは、 --nproc_per_nodeを調整することにより、複数のGPUで並列化できます。通常、2番目のコマンドは1〜3分かかりますが、最初のコマンドは構成に応じて数時間かかる場合があります。オプションの完全なリストについてはpython fid.py --helpを参照してください。
FIDの数値は、異なるランダムシードによって異なり、画像の数に非常に敏感であることに注意してください。デフォルトでは、 fid.py常に50,000の生成された画像を使用します。より少ない画像を提供するとエラーが発生しますが、より多くを提供するとランダムサブセットが使用されます。ランダムバリエーションの影響を減らす--seeds=50000-99999に、異なる種子で複数回計算--seeds=0-49999繰り返すことをお勧めし--seeds=100000-149999 。私たちの論文では、各FIDを3回計算し、最小を報告しました。
また、モデルが元々訓練されていたのと同じデータセットと生成された画像を比較することが重要であることに注意してください。評価を容易にするために、事前に訓練されたモデルに対応する正確な参照統計を提供します。
ImagENETの場合、リンゴからアプリルの比較を有効にするための2セットの参照統計を提供します。EDMモデル（ edm-imagenet-64x64-cond-adm.pkl ）を評価するときは、 imagenet-64x64-baseline.npz imagenet-64x64.npz使用する必要があります。 baseline-imagenet-64x64-cond-adm.pkl ）;後者はもともと、わずかに異なるトレーニングデータを使用して、DhariwalとNicholによって訓練されました。
次のように、独自のデータセットの参照統計を計算できます。
python fid.py ref --data=datasets/my-dataset.zip --dest=fid-refs/my-dataset.npzデータセットは、スタイルガンと同じ形式で保存されます。非圧縮PNGファイルとラベル用のメタデータファイルdataset.jsonを含む非圧縮ZIPアーカイブ。カスタムデータセットは、画像を含むフォルダーから作成できます。詳細については、 python dataset_tool.py --help参照してください。
CIFAR-10： CIFAR-10 Pythonバージョンをダウンロードして、Zip Archiveに変換します。
python dataset_tool.py --source=downloads/cifar10/cifar-10-python.tar.gz
--dest=datasets/cifar10-32x32.zip
python fid.py ref --data=datasets/cifar10-32x32.zip --dest=fid-refs/cifar10-32x32.npzFFHQ： Flickr-Faces-HQデータセットを1024x1024画像としてダウンロードし、64x64解像度でZIPアーカイブに変換します。
python dataset_tool.py --source=downloads/ffhq/images1024x1024
--dest=datasets/ffhq-64x64.zip --resolution=64x64
python fid.py ref --data=datasets/ffhq-64x64.zip --dest=fid-refs/ffhq-64x64.npz AFHQV2：更新された動物FACES-HQデータセット（ afhq-v2-dataset ）をダウンロードし、64x64解像度でZIPアーカイブに変換します。
python dataset_tool.py --source=downloads/afhqv2
--dest=datasets/afhqv2-64x64.zip --resolution=64x64
python fid.py ref --data=datasets/afhqv2-64x64.zip --dest=fid-refs/afhqv2-64x64.npzImagenet： ImagENetオブジェクトのローカリゼーションチャレンジをダウンロードし、64x64解像度でzipアーカイブに変換します。
python dataset_tool.py --source=downloads/imagenet/ILSVRC/Data/CLS-LOC/train
--dest=datasets/imagenet-64x64.zip --resolution=64x64 --transform=center-crop
python fid.py ref --data=datasets/imagenet-64x64.zip --dest=fid-refs/imagenet-64x64.npztrain.pyを使用して新しいモデルをトレーニングできます。例えば：
# Train DDPM++ model for class-conditional CIFAR-10 using 8 GPUs
torchrun --standalone --nproc_per_node=8 train.py --outdir=training-runs
--data=datasets/cifar10-32x32.zip --cond=1 --arch=ddpmpp上記の例では、8 GPU（ --nproc_per_nodeで制御）で均等に分割される512画像のデフォルトバッチサイズ（ --batchで制御）を使用して、GPUあたり64画像を生成します。大規模なモデルのトレーニングは、GPUメモリがなくなる場合があります。これを回避する最良の方法は、GPUごとのバッチサイズ--batch-gpu=32を制限することです。これにより、勾配蓄積が使用され、完全なGPUバッチを使用するのと同じ結果が得られます。オプションの完全なリストについては、 python train.py --help参照してください。
各トレーニングの結果は、たとえばtraining-runs/00000-cifar10-cond-ddpmpp-edm-gpus8-batch64-fp32など、新しく作成されたディレクトリに保存されます。トレーニングループは、ネットワークスナップショット（ network-snapshot-*.pkl ）とトレーニング状態（ training-state-*.pt ）を定期的にエクスポートし、 --snapおよび--dumpで制御）。ネットワークスナップショットを使用してgenerate.pyを使用して画像を生成でき、トレーニング状態を使用してトレーニングを後で再開できます（ --resume ）。その他の有用な情報は、 log.txtおよびstats.jsonlに記録されています。トレーニングの収束stats.jsonl監視するには、 generate.pyおよびfid.py "Loss/loss"使用して、 network-snapshot-*.pkl
次の表には、事前に訓練されたモデルを取得するために使用した正確なトレーニング構成を示します。
| モデル | GPU | 時間 | オプション |
|---|---|---|---|
| CIFAR10‑32X32 ‑ COND ‑ VP | 8xv100 | 〜2日 | --cond=1 --arch=ddpmpp |
| CIFAR10‑32X32 ‑ COND -ve | 8xv100 | 〜2日 | --cond=1 --arch=ncsnpp |
| CIFAR10‑32X32 -UNCOND ‑ VP | 8xv100 | 〜2日 | --cond=0 --arch=ddpmpp |
| CIFAR10‑32X32 -UNCOND -ve | 8xv100 | 〜2日 | --cond=0 --arch=ncsnpp |
| ffhq ‑ 64x64 ‑ uncond ‑ vp | 8xv100 | 〜4日 | --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 |
| ffhq ‑ 64x64 ‑ uncond ‑ ve | 8xv100 | 〜4日 | --cond=0 --arch=ncsnpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.05 --augment=0.15 |
| AFHQV2‑64X64 -UNCOND ‑ VP | 8xv100 | 〜4日 | --cond=0 --arch=ddpmpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.25 --augment=0.15 |
| afhqv2‑64x64 ‑ uncond -ve | 8xv100 | 〜4日 | --cond=0 --arch=ncsnpp --batch=256 --cres=1,2,2,2 --lr=2e-4 --dropout=0.25 --augment=0.15 |
| Imagenet ‑ 64x64 ‑ cond ‑ adm | 32XA100 | 〜13日 | --cond=1 --arch=adm --duration=2500 --batch=4096 --lr=1e-4 --ema=50 --dropout=0.10 --augment=0 --fp16=1 --ls=100 --tick=200 |
Imagenet-64の場合、4つのNVIDIA DGX A100ノードでトレーニングを実行しました。それぞれに80 GBのメモリを備えた8つのアンペアGPUが含まれています。 GPUメモリ要件を削減するには、より多くのGPUでモデルをトレーニングするか、 --batch-gpuでGPUごとのバッチサイズを制限することをお勧めします。マルチノードトレーニングをセットアップするには、Torchrunドキュメントを参照してください。
Copyright©2022、Nvidia Corporation＆Affiliates。無断転載を禁じます。
ソースコードや事前に訓練されたモデルを含むすべての資料は、Creative Commons Attribution-Commercial-Sharealike 4.0 International Licenseの下でライセンスされています。
baseline-cifar10-32x32-uncond-vp.pklおよびbaseline-cifar10-32x32-uncond-ve.pklは、ヤンソング、Jascha Sohl-Dickstein、Diederik P. Kingma、Abhishek Kumar、Stefano Ermon、Ben Pooleの事前に訓練されたモデルから派生しています。モデルはもともと、Apache 2.0ライセンスの下で共有されていました。
baseline-imagenet-64x64-cond-adm.pklは、Prafulla DhariwalとAlex Nicholによって事前に訓練されたモデルから派生しています。モデルはもともとMITライセンスの下で共有されていました。
imagenet-64x64-baseline.npz Prafulla DhariwalとAlex Nicholによる事前計算された参照統計から派生しています。統計はもともとMITライセンスの下で共有されていました。
@inproceedings{Karras2022edm,
author = {Tero Karras and Miika Aittala and Timo Aila and Samuli Laine},
title = {Elucidating the Design Space of Diffusion-Based Generative Models},
booktitle = {Proc. NeurIPS},
year = {2022}
}
これは研究リファレンスの実装であり、1回限りのコードドロップとして扱われます。そのため、プルリクエストの形で外部コードの貢献を受け入れません。
Jaako Lehtinen、Ming-Yu Liu、TuomasKynkäänniemi、Axel Sauer、Arash Vahdat、およびJanne Hellstenの議論とコメント、およびTero Kuosmanen、Samuel Klenberg、Janne HellstenのCultute Infrastrutureを維持してくれたことに感謝します。
