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 หมายเหตุ: อิมเมจนักเทียบท่าต้องการไดรเวอร์ NVIDIA RELEEST r520 หรือใหม่กว่า
การเรียกใช้ docker run อาจดูน่ากลัวดังนั้นลองแกะเนื้อหาที่นี่:
--gpus all -it --rm --user $(id -u):$(id -g) : เมื่อเปิดใช้งาน GPU ทั้งหมดให้เรียกใช้เซสชันการโต้ตอบกับ UID/GID ของผู้ใช้ปัจจุบันเพื่อหลีกเลี่ยงการเขียนไฟล์เป็นรูท-v `pwd`:/scratch --workdir /scratch : เมานต์ปัจจุบัน Running Dir (เช่นด้านบนของ repo git นี้บนเครื่องโฮสต์ของคุณ) ถึง /scratch ในคอนเทนเนอร์และใช้เป็น dir ที่ทำงานในปัจจุบัน-e HOME=/scratch : ระบุตำแหน่งที่จะแคชไฟล์ชั่วคราว หมายเหตุ: หากคุณต้องการการควบคุมที่ละเอียดยิ่งขึ้นคุณสามารถตั้งค่า DNNLIB_CACHE_DIR (สำหรับแคชดาวน์โหลดรุ่นที่ผ่านการฝึกอบรมมาแล้ว) คุณต้องการให้แคชเหล่านี้อยู่ในปริมาณที่ถาวรเพื่อให้เนื้อหาของพวกเขาถูกเก็บไว้ในการเรียกใช้ docker run หลายรายการ เรามีโมเดลที่ผ่านการฝึกอบรมมาล่วงหน้าสำหรับการกำหนดค่าการฝึกอบรมที่เราเสนอ (config f) รวมถึงการกำหนดค่าพื้นฐาน (config 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 การสร้างภาพจำนวนมากอาจใช้เวลานาน เวิร์กโหลดสามารถแจกจ่ายผ่าน GPU หลายตัวโดยการเรียกใช้คำสั่งด้านบนโดยใช้ 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) สำหรับรุ่นและตัวอย่างที่กำหนดให้สร้างภาพสุ่ม 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 คำสั่งทั้งสองข้างต้นสามารถขนานกับ GPU หลายตัวโดยการปรับ --nproc_per_node โดยทั่วไปแล้วคำสั่งที่สองใช้เวลา 1-3 นาทีในทางปฏิบัติ แต่บางครั้งคำแรกอาจใช้เวลาหลายชั่วโมงขึ้นอยู่กับการกำหนดค่า ดู python fid.py --help สำหรับรายการตัวเลือกทั้งหมด
โปรดทราบว่าค่าตัวเลขของ FID นั้นแตกต่างกันไปตามเมล็ดสุ่มที่แตกต่างกันและมีความไวสูงต่อจำนวนภาพ โดยค่าเริ่มต้น fid.py จะใช้ภาพที่สร้างขึ้น 50,000 ภาพเสมอ การให้ภาพน้อยลงจะส่งผลให้เกิดข้อผิดพลาดในขณะที่การให้บริการเพิ่มเติมจะใช้ชุดย่อยแบบสุ่ม เพื่อลดผลกระทบของการเปลี่ยนแปลงแบบสุ่มเราขอแนะนำให้ทำการคำนวณซ้ำหลายครั้งด้วยเมล็ดที่แตกต่างกันเช่น --seeds=0-49999 , --seeds=50000-99999 และ --seeds=100000-149999 ในบทความของเราเราคำนวณ FID แต่ละครั้งสามครั้งและรายงานขั้นต่ำ
นอกจากนี้โปรดทราบว่ามันเป็นสิ่งสำคัญที่จะเปรียบเทียบภาพที่สร้างขึ้นกับชุดข้อมูลเดียวกับที่รุ่นได้รับการฝึกฝนด้วย เพื่ออำนวยความสะดวกในการประเมินผลเราจัดทำสถิติอ้างอิงที่แน่นอนซึ่งสอดคล้องกับโมเดลที่ผ่านการฝึกอบรมมาก่อน:
สำหรับ Imagenet เรามีสถิติอ้างอิงสองชุดเพื่อเปิดใช้งานการเปรียบเทียบแอปเปิ้ลกับแอปเปิ้ล: imagenet-64x64.npz ควรใช้เมื่อประเมินแบบจำลอง EDM ( edm-imagenet-64x64-cond-adm.pkl imagenet-64x64-baseline.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: คลังเก็บซิปที่ไม่มีการบีบอัดที่มีไฟล์ PNG ที่ไม่บีบอัดและ dataset.json ไฟล์ข้อมูลเมตา ชุดข้อมูลที่กำหนดเองสามารถสร้างได้จากโฟลเดอร์ที่มีรูปภาพ ดู 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 และแปลงเป็น zip เก็บถาวรที่ความละเอียด 64x64:
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 เก็บถาวรที่ความละเอียด 64x64:
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 Object Localization Challenge และแปลงเป็น Zip Archive ที่ความละเอียด 64x64:
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 GPU (ควบคุมโดย --nproc_per_node ) เพื่อให้ได้ 64 ภาพต่อ GPU การฝึกอบรมแบบจำลองขนาดใหญ่อาจหมดหน่วยความจำ 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 ในการตรวจสอบการบรรจบกันของการฝึกอบรมเราขอแนะนำให้ดูการสูญเสียการฝึกอบรม ( "Loss/loss" ใน stats.jsonl ) รวมถึงการประเมิน FID เป็นระยะสำหรับ network-snapshot-*.pkl โดยใช้ generate.py และ fid.py
ตารางต่อไปนี้แสดงการกำหนดค่าการฝึกอบรมที่แน่นอนที่เราใช้เพื่อรับโมเดลที่ผ่านการฝึกอบรมมาก่อน:
| แบบอย่าง | GPUs | เวลา | ตัวเลือก |
|---|---|---|---|
| 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 - ADM - 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 แอมป์ GPU ที่มีหน่วยความจำ 80 GB เพื่อลดความต้องการหน่วยความจำ GPU เราขอแนะนำให้ฝึกอบรมแบบจำลองด้วย GPU มากขึ้นหรือ จำกัด ขนาดแบทช์ต่อ GPU ด้วย --batch-gpu หากต้องการตั้งค่าการฝึกอบรมแบบหลายโหนดโปรดปรึกษาเอกสารประกอบ Torchrun
ลิขสิทธิ์© 2022, Nvidia Corporation & Affiliates สงวนลิขสิทธิ์
วัสดุทั้งหมดรวมถึงซอร์สโค้ดและโมเดลที่ผ่านการฝึกอบรมมาก่อนได้รับอนุญาตภายใต้ใบอนุญาต International Creative Commons Commons-Noncommercial-Sharealike 4.0
baseline-cifar10-32x32-uncond-vp.pkl และ baseline-cifar10-32x32-uncond-ve.pkl มาจากแบบจำลองที่ผ่านการฝึกอบรมมาก่อนโดย Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek 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 สถิติเดิมใช้ร่วมกันภายใต้ใบอนุญาต 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}
}
นี่คือการดำเนินการอ้างอิงการวิจัยและถือว่าเป็นการลดลงของรหัสครั้งเดียว ดังนั้นเราจึงไม่ยอมรับการมีส่วนร่วมของรหัสภายนอกในรูปแบบของคำขอดึง
เราขอขอบคุณ Jaakko Lehtinen, Ming-yu Liu, Tuomas Kynkänniemi, Axel Sauer, Arash Vahdat และ Janne Hellsten สำหรับการอภิปรายและความคิดเห็นและ Tero Kuosmanen, Samuel Klenberg และ Janne Hellsten
