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 الخارجية ، فقط تأكد من أن torch_utils و dnnlib يتم تصنيعهما من خلال PYTHONPATH .
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 شاقًا ، لذلك دعونا نتفكك محتوياته هنا:
--gpus all -it --rm --user $(id -u):$(id -g) : مع تمكين جميع وحدات معالجة الرسومات ، قم بتشغيل جلسة تفاعلية مع UID/GID للمستخدم الحالي لتجنب ملفات Docker كجذر.-v `pwd`:/scratch --workdir /scratch : mount current running dir (على سبيل المثال ، الجزء العلوي من هذا الريبو git على جهاز المضيف الخاص بك) إلى /scratch في الحاوية واستخدامه كدليل العمل الحالي.-e HOME=/scratch : حدد مكان ذاكرة التخزين المؤقتة للملفات المؤقتة. ملاحظة: إذا كنت تريد المزيد من التحكم في الحبيبات ، فيمكنك بدلاً من ذلك تعيين DNNLIB_CACHE_DIR (لتنزيل تنزيل النموذج الذي تم تدريبه مسبقًا). تريد أن تقيم هذه الأنباء ذاكرة التخزين المؤقت على أحجام مستمرة بحيث يتم الاحتفاظ بمحتوياتها عبر دعوات متعددة docker run . نحن نقدم نماذج مدربة مسبقًا لتكوين التدريب المقترح (التكوين F) وكذلك تكوين خط الأساس (التكوين أ):
لإنشاء مجموعة من الصور باستخدام نموذج معين وأخذ عينات ، قم بتشغيل:
# 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 :
# 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échet (FID) لنموذج معين وأخذ عينات ، قم أولاً بإنشاء 50000 صورة عشوائية ثم قارنها مقابل إحصائيات مرجعية مجموعة البيانات باستخدام 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 . يستغرق الأمر الثاني عادةً 1-3 دقائق في الممارسة ، ولكن يمكن للأول الأول في بعض الأحيان أن يستغرق عدة ساعات ، اعتمادًا على التكوين. انظر python fid.py --help للحصول على القائمة الكاملة للخيارات.
لاحظ أن القيمة العددية لـ FID تختلف عبر البذور العشوائية المختلفة وهي حساسة للغاية لعدد الصور. بشكل افتراضي ، سوف يستخدم fid.py دائمًا 50000 صورة تم إنشاؤها ؛ سيؤدي توفير عدد أقل من الصور إلى خطأ ، في حين أن توفير المزيد سيستخدم مجموعة فرعية عشوائية. لتقليل تأثير التباين العشوائي ، نوصي بتكرار الحساب عدة مرات مع بذور مختلفة ، على سبيل المثال ، --seeds=0-49999 ، --seeds=50000-99999 ، و- --seeds=100000-149999 . في ورقتنا ، قمنا بحساب كل FID ثلاث مرات وأبلغنا الحد الأدنى.
لاحظ أيضًا أنه من المهم مقارنة الصور التي تم إنشاؤها مقابل نفس مجموعة البيانات التي تم تدريبها في الأصل. لتسهيل التقييم ، نقدم الإحصاءات المرجعية الدقيقة التي تتوافق مع نماذجنا التي تم تدريبنا قبل التدريب:
بالنسبة edm-imagenet-64x64-cond-adm.pkl ImageNet ، نقدم مجموعتين من الإحصائيات المرجعية لتمكين مقارنة 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 يتم تخزين مجموعات البيانات بنفس التنسيق كما في Stylegan: أرشيفات zip غير مضغوطة تحتوي على ملفات PNG غير مضغوطة ومجموعة dataset.json ملفات تعريف البيانات الوصفية. يمكن إنشاء مجموعات بيانات مخصصة من مجلد يحتوي على صور ؛ انظر python dataset_tool.py --help لمزيد من المعلومات.
CIFAR-10: قم بتنزيل نسخة CIFAR-10 Python وتحويلها إلى أرشيف ZIP:
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 كصور 1024 × 1024 وتحويلها إلى أرشيف zip بدقة 64 × 64:
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: قم بتنزيل مجموعة بيانات Animal Faces-HQ المحدثة ( afhq-v2-dataset ) وتحويلها إلى أرشيف ZIP بدقة 64 × 64:
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 وتحويله إلى أرشيف ZIP بدقة 64 × 64:
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.npz يمكنك تدريب طرز جديدة باستخدام train.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 يستخدم المثال أعلاه حجم الدفعة الافتراضية لـ 512 صورة (يتم التحكم فيها بواسطة --batch ) مقسمة بالتساوي بين 8 وحدات معالجة الرسومات (التي يتم التحكم فيها بواسطة --nproc_per_node ) لإعطاء 64 صورة لكل وحدة معالجة الرسومات. قد ينفد تدريب النماذج الكبيرة من ذاكرة GPU ؛ أفضل طريقة لتجنب ذلك هي الحد من حجم دفعة GPU ، على سبيل المثال ، --batch-gpu=32 . هذا يستخدم تراكم التدرج لتحقيق نفس النتائج مثل استخدام دفعات كاملة لكل GPU. انظر python train.py --help للحصول على القائمة الكاملة للخيارات.
يتم حفظ نتائج كل تشغيل تدريب إلى دليل تم إنشاؤه حديثًا ، على سبيل المثال training-runs/00000-cifar10-cond-ddpmpp-edm-gpus8-batch64-fp32 . --dump حلقة التدريب لقطات الشبكة ( network-snapshot-*.pkl --snap وحالات التدريب ( training-state-*.pt . يمكن استخدام لقطات الشبكة لإنشاء صور باستخدام generate.py ، ويمكن استخدام حالات التدريب لاستئناف التدريب لاحقًا ( --resume ). يتم تسجيل معلومات مفيدة أخرى في log.txt و stats.jsonl . لمراقبة تقارب التدريب ، نوصي generate.py إلى فقدان التدريب ( "Loss/loss" في stats.jsonl ) وكذلك تقييم FID بشكل دوري fid.py network-snapshot-*.pkl
يسرد الجدول التالي تكوينات التدريب الدقيقة التي استخدمناها للحصول على نماذجنا التي تم تدريبنا مسبقًا:
| نموذج | وحدات معالجة الرسومات | وقت | خيارات |
|---|---|---|---|
| CIFAR10‑32X32 - Cond - VP | 8xv100 | ~ 2 أيام | --cond=1 --arch=ddpmpp |
| CIFAR10‑32X32 - Cond - | 8xv100 | ~ 2 أيام | --cond=1 --arch=ncsnpp |
| CIFAR10‑32X32 - UNCOND - VP | 8xv100 | ~ 2 أيام | --cond=0 --arch=ddpmpp |
| Cifar10‑32x32 - uncond - | 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 - | 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 - | 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 ، قمنا بتدريبنا على أربعة عقد NVIDIA DGX A100 ، كل منها يحتوي على 8 أمبير مع معالجة الرسومات مع 80 جيجابايت من الذاكرة. لتقليل متطلبات ذاكرة GPU ، نوصي إما بتدريب النموذج مع مزيد من وحدات معالجة الرسومات أو الحد من حجم دفعة GPU مع-- --batch-gpu . لإعداد تدريب متعدد العقدة ، يرجى الرجوع إلى وثائق Torchrun.
حقوق الطبع والنشر © 2022 ، Nvidia Corporation & Affiliates. جميع الحقوق محفوظة.
يتم ترخيص جميع المواد ، بما في ذلك التعليمات البرمجية المصدرية والنماذج المدربة مسبقًا ، بموجب ترخيص Creative Commons Noncommercial-Sharealike 4.0 الدولي.
يتم اشتقاق baseline-cifar10-32x32-uncond-vp.pkl و baseline-cifar10-32x32-uncond-ve.pkl من النماذج التي تم تدريبها مسبقًا من قبل Yang Song و Jascha Sohl-Dickstein و Dieterik P. Kingma و Abhishekek Kumar و Stefano ermon. تمت مشاركة النماذج في الأصل بموجب ترخيص Apache 2.0.
تم اشتقاق baseline-imagenet-64x64-cond-adm.pkl من النموذج الذي تم تدريبه مسبقًا بواسطة Prafulla Dhariwal و Alex Nichol. تمت مشاركة النموذج في الأصل بموجب ترخيص MIT.
imagenet-64x64-baseline.npz مشتق من الإحصاءات المرجعية المسبقة من قبل Prafulla Dhariwal و Alex Nichol. تمت مشاركة الإحصاءات في الأصل بموجب ترخيص معهد ماساتشوستس للتكنولوجيا.
@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}
}
هذا هو تطبيق مرجعي للبحث ويتم التعامل معه على أنه قطرة رمز لمرة واحدة. على هذا النحو ، لا نقبل مساهمات الكود الخارجية في شكل طلبات سحب.
نشكر Jaakko Lehtinen و Ming-Yu Liu و Tuomas Kynkäänniemi و Axel Sauer و Arash Vahdat و Janne Hellsten على المناقشات والتعليقات ، وتيرو كوسمانين ، صموئيل كلنبرغ ، وجان هيلستن لمحافظة كومبوتنا.
