إسماعيل بن علا

ImageGen

شبكات GAN ونماذج الانتشار خلف حلقة تدريب واحدة

الذكاء الاصطناعي التوليدي · يونيو 2024 · قيد الإنجاز

DCGAN وGAN من نوع ResNet وشبكة U-Net للانتشار، تتدرب وتحفظ نقاطها وتولّد العينات بالطريقة نفسها، من MNIST على CPU إلى وجوه الحيوانات على GPU.

التقنياتPyTorchtorchvisionHugging Face Datasetstorchrunpytest

عرض الكود على GitHub
ImageGen — نظرة عامة

نظرة عامة

أردتُ توليد وجوه حيوانات وفهم عائلتي النماذج القادرتين على ذلك: شبكات GAN، حيث يُدرَّب مولّد ومميِّز أحدهما ضد الآخر، ونماذج الانتشار (Diffusion Models)، حيث تتعلم شبكة U-Net إزالة قليل من الضجيج ويكون التوليد إزالةً للضجيج من ضجيج خالص خطوة بعد خطوة. يضع ImageGen شبكة DCGAN وشبكة GAN من نوع ResNet على نمط SNGAN/BigGAN وشبكة U-Net لنموذج DDPM خلف حلقة تدريب واحدة وأداة توليد واحدة وصيغة نقاط حفظ واحدة، بـ PyTorch خالص.

يمكن أن تكون النماذج الثلاثة شرطية حسب الفئة («ارسم قطًا»)، وتحتفظ بنسخة EMA من الأوزان للتوليد، وتستأنف من نقطة حفظ ذرّية تخزّن المُحسِّنات وEMA وحالة المولّد العشوائي والإعدادات، وتعمل على CPU أو GPU واحد أو عدة GPU أو عدة أجهزة عبر torchrun. وتغطي الإعدادات الجاهزة MNIST وCIFAR-10 وAFHQ (16,130 وجه حيوان في ثلاث فئات: قط، كلب، بري) بدقة 64 و128 px، مع أمر تجهيز يجلب AFHQ تدفقيًا من Hugging Face Hub إلى مجلدات فئات بأحجام معدّلة. ويمكن حساب FID أثناء التدريب، ويتتبّع best.pt أدنى قيمة.

بدأ المشروع دفترَ Colab فيه GAN مبني على ResBlock يعمل على ملف zip لصور حيوانات؛ كان يدرّب على متجهات الضجيج الـ 32 نفسها في كل خطوة وفيه عدة أخطاء أخرى، كلها مُدرجة ومُصلحة في الحزمة. وحتى الآن تم التحقق من كل معمارية وكل دالة خسارة لـ GAN على MNIST باستخدام CPU، بما في ذلك التركيبات التي لم تنجح. أما التشغيلات الكاملة على GPU مع CIFAR-10 وAFHQ وأرقام FID ونقاط الحفظ المنشورة فهي الخطوة التالية.

السؤال

هل تستطيع قاعدة شيفرة واحدة أن تدرّب عائلتي GAN ونموذج انتشار بالحلقة نفسها ونقطة الحفظ نفسها وواجهة التوليد نفسها، وأن تتوسع من تشغيل تحقّق على CPU مع MNIST إلى وجوه الحيوانات AFHQ على عدة GPU؟

أهم الميزات

كيف يعمل

  1. بيانات ضمن [-1, 1]

    يُعاد تحجيم الصور على الضلع الأقصر، وتُقص من المركز، وتُحوَّل إلى المجال [-1, 1]، كي يتشارك مخرج tanh في المولّدات وضجيج الانتشار مقياس البيانات نفسه. وتصبح المجلدات الفرعية فئات.

  2. خطوة GAN

    تُكرَّر d_steps مرة: توليد صور مزيفة دون تدرجات نحو G، وتقييم الحقيقية والمزيفة، وإضافة عقوبة التدرج أو R1 عند الطلب، وتحديث D. ثم تجميد D، وتوليد صور مزيفة جديدة، وتحديث G عبر D، وتحديث نسخة EMA من G المستخدمة في التوليد.

  3. خطوة انتشار

    إضافة ضجيج عند خطوة زمنية عشوائية من الجدولة الجيبية والانحدار نحو ε أو v بمتوسط مربع الخطأ؛ ويجري التوليد بـ DDIM (أو بأسلوب DDPM مع η = 1) انطلاقًا من ضجيج خالص، مع classifier-free guidance للنماذج الشرطية.

  4. حلقة واحدة للعائلتين

    لا تستدعي الحلقة إلا task.train_step()، وتكتب last.pt وشبكات معاينة من ضجيج ثابت كل بضع خطوات، وتسجّل المقاييس في ملف JSONL، وعند الطلب تحسب FID وتحتفظ بـ best.pt.

  5. التوسّع

    تحت torchrun تحسب كل عملية تدرجاتها الخاصة، ويأخذ all-reduce واحد متوسطها قبل كل خطوة للمُحسِّن؛ ولا تكتب الملفات إلا العملية ذات الرتبة 0. ويشغّل اختبارٌ عمليتين على CPU ويتحقق من انتهائهما بأوزان متطابقة.

أمثلة

GAN شرطي من نوع ResNet (hinge loss وspectral norm وself-attention ومميِّز إسقاطي، 32 قناة، خطوتان لـ D) بعد 3,000 خطوة على CPU: أفضل نتيجة لـ GAN، كل صف فئة رقم واحدة، وكانت الصفوف تتبع الوسوم منذ الخطوة 500.
GAN شرطي من نوع ResNet (hinge loss وspectral norm وself-attention ومميِّز إسقاطي، 32 قناة، خطوتان لـ D) بعد 3,000 خطوة على CPU: أفضل نتيجة لـ GAN، كل صف فئة رقم واحدة، وكانت الصفوف تتبع الوسوم منذ الخطوة 500.
DCGAN شرطي بدالة الخسارة غير المشبعة بعد 4,000 خطوة على CPU: الرقم الصحيح في كل صف منذ الخطوة 500، مع d_loss حول 0.3 وg_loss حول 3 طوال التدريب.
DCGAN شرطي بدالة الخسارة غير المشبعة بعد 4,000 خطوة على CPU: الرقم الصحيح في كل صف منذ الخطوة 500، مع d_loss حول 0.3 وg_loss حول 3 طوال التدريب.
شبكة U-Net للانتشار بتنبؤ v بعد 2,500 خطوة (16 قناة، guidance 2): خلفيات سوداء نظيفة والرقم الصحيح في كل صف؛ وانخفضت الخسارة من 1.0 إلى 0.095.
شبكة U-Net للانتشار بتنبؤ v بعد 2,500 خطوة (16 قناة، guidance 2): خلفيات سوداء نظيفة والرقم الصحيح في كل صف؛ وانخفضت الخسارة من 1.0 إلى 0.095.
شبكة U-Net نفسها بتنبؤ ε: الأرقام صحيحة لكن كثيرًا من العينات تحتفظ بخلفيات رمادية مرقّطة، لأن استرجاع x0 من ε يضخّم الأخطاء بنحو 70x في الخطوات الأكثر ضجيجًا. ولهذا انتقلت الإعدادات الجاهزة إلى تنبؤ v.
شبكة U-Net نفسها بتنبؤ ε: الأرقام صحيحة لكن كثيرًا من العينات تحتفظ بخلفيات رمادية مرقّطة، لأن استرجاع x0 من ε يضخّم الأخطاء بنحو 70x في الخطوات الأكثر ضجيجًا. ولهذا انتقلت الإعدادات الجاهزة إلى تنبؤ v.

المدخل

imagegen-train --preset mnist-dcgan --conditional --steps 4000 --output-dir checkpoints/mnist-dcgan imagegen-sample --checkpoint checkpoints/mnist-dcgan/last.pt --class 7 --num 16 --out sevens.png

المخرج

checkpoints/mnist-dcgan/last.pt, samples/step_*.png and metrics.jsonl, then sevens.png: 16 generated sevens from the EMA generator

ثنائي البدء السريع: تدريب DCGAN شرطي على MNIST، ثم طلب فئة واحدة من نقطة الحفظ. تخزّن نقطة الحفظ إعداداتها، فلا يحتاج التوليد إلى أي خيارات خاصة بالمعمارية.

التقنيات المستخدمة

ما تعلمته

الحدود

الكلمات المفتاحية

توليد الصور · GAN · DCGAN · SNGAN · BigGAN · التطبيع الطيفي · hinge loss · المميِّز الإسقاطي · DiffAugment · نماذج الانتشار · DDPM · DDIM · تنبؤ v · classifier-free guidance · EMA · AFHQ · CIFAR-10 · MNIST · التدريب الموزّع · PyTorch