FlaxDiff
1.0.0
يتم دعم هذا المشروع جزئيًا بواسطة Cloud Google TPU Research Cloud. أود أن أشكر فريق Google Cloud TPU على تزويدني بالموارد لتدريب النماذج الكبرى في النصوص في الإعدادات الموزعة متعددة المضيفات.
في السنوات الأخيرة ، أحدثت النماذج متعددة الخطوات القائمة على الانتشار والنتيجة ثورة في مجال الذكاء الاصطناعي. ومع ذلك ، فقد أصبح آخر الأبحاث في هذا المجال كثافة في الرياضيات ، مما يجعل من الصعب فهم كيفية عمل نماذج الانتشار الحديثة وتوليد مثل هذه الصور المثيرة للإعجاب. يمكن أن يكون تكرار هذا البحث في الكود أمرًا شاقًا.
FlaxDiff هي مكتبة من الأدوات (الجدولين ، أخذ العينات ، النماذج ، إلخ) المصممة وتنفيذها بطريقة سهلة الفهم. ينصب التركيز على قابلية الفهم وقابلية القراءة على الأداء. لقد بدأت هذا المشروع كهواية للتعرف على الكتان وجاكس والتعرف على الانتشار وأحدث الأبحاث في الذكاء الاصطناعي.
لقد بدأت في البداية هذا المشروع في Keras ، على دراية بـ TensorFlow 2.0 ، ولكن انتقلت إلى الكتان ، مدعوم من Jax ، لأدائه وسهولة الاستخدام. يتم توفير دفاتر الملاحظات والموديلات القديمة ، بما في ذلك نماذج الكتان الأولى ، أيضًا.
دفتر Noteboor Diffusion_flax_linen.ipynb هو مساحة العمل الرئيسية الخاصة بي للتجارب. يتم تحميل العديد من نقاط التفتيش على المجلد pretrained مع نسخة من دفتر ملاحظات العمل المرتبط بكل نقطة تفتيش. قد تحتاج إلى نسخ دفتر الملاحظات إلى جذر العمل حتى يعمل بشكل صحيح.
في مجلد example notebooks ، ستجد أجهزة الكمبيوتر المحمولة الشاملة لتقنيات الانتشار المختلفة ، مكتوبة بالكامل من نقطة الصفر وهي مستقلة عن مكتبة FlaxDiff. يتضمن كل دفتر ملاحظات تفسيرات مفصلة للرياضيات والمفاهيم الأساسية ، مما يجعلها موارد لا تقدر بثمن للتعلم وفهم نماذج الانتشار.
شرح الانتشار (رابط NBViewer) (الرابط المحلي)
EDM (توضيح مساحة تصميم النماذج التوليدية القائمة على الانتشار)
تهدف دفاتر الملاحظات هذه إلى توفير دليل سهل الفهم وخطوة بخطوة لنماذج وتقنيات الانتشار المختلفة. تم تصميمها لتكون صديقة للمبتدئين ، وبالتالي على الرغم من أنها قد لا تلتزم بالركضات والتطبيقات الدقيقة للأوراق الأصلية لجعلها أكثر فهمًا وقابلة للتعميم ، فقد بذلت قصارى جهدي للحفاظ عليها دقيقة قدر الإمكان. إذا وجدت أي أخطاء أو لديك أي اقتراحات ، فلا تتردد في فتح مشكلة أو طلب سحب.
برنامج تدريب متوازي بيانات متعدد المضيفات في جاكس
مرافق TPU لجعل الحياة أسهل
عملت كباحث في التعلم الآلي في Hyperverge من 2019-2021 ، مع التركيز على رؤية الكمبيوتر ، وخاصة الوجه لمكافحة الإرشاد والكشف عن الوجه والتعرف عليها. منذ التحول إلى وظيفتي الحالية في عام 2021 ، لم أشارك في الكثير من أعمال البحث والتطوير ، مما دفعني إلى بدء مشروع الحيوانات الأليفة هذا لإعادة النظر في الأساسيات وإعادة تعلمها والتعرف على أحدث الأحد. يتضمن دوري الحالي في المقام الأول هندسة نظام Golang مع بعض أعمال ML المطبقة التي تم رشها للتو. لذلك ، قد يعكس الرمز رحلتي التعليمية. أرجوك سامح أي أخطاء وفتح مشكلة لإخبارنا.
أيضًا ، قد يتم إنشاء عدد قليل من النص بمساعدة Github Copilot ، لذا يرجى تعذر أي أخطاء في النص.
تم تنفيذه في flaxdiff.schedulers :
flaxdiff.schedulers.LinearNoiseSchedule ): جدولة منفصلة مع المعلمة بيتا.flaxdiff.schedulers.CosineNoiseSchedule ): جدولة منفصلة بالبيوتا.flaxdiff.schedulers.ExpNoiseSchedule ): جدولة منفصلة مع المعلمة بيتا.flaxdiff.schedulers.CosineContinuousNoiseScheduler ): جدولة مستمرة.flaxdiff.schedulers.CosineGeneralNoiseScheduler ): جدولة جسم تجضية مستمرة Sigma.flaxdiff.schedulers.KarrasVENoiseScheduler ): جدولة مستمرة مع المعلمة التي اقترحها Karras et al. 2022 ، الأنسب للاستدلال.flaxdiff.schedulers.EDMNoiseScheduler ): جدولة مستمرة معجزة على أساس نموذج الانتشار الأسي (EDM) ، الأنسب للتدريب مع KarraskarrasveniseCheduler. تم تنفيذه في flaxdiff.predictors :
flaxdiff.predictors.EpsilonPredictor ): يتوقع الضوضاء في البيانات.flaxdiff.predictors.X0Predictor ): يتوقع البيانات الأصلية من البيانات الصاخبة.flaxdiff.predictors.VPredictor ): يتنبأ بمزيج خطي من البيانات والضوضاء ، شائعة الاستخدام في EDM.flaxdiff.predictors.KarrasEDMPredictor ): مؤشر عام لـ EDM ، يدمج المعلمات المختلفة. تم تنفيذه في flaxdiff.samplers :
flaxdiff.samplers.DDPMSampler ): ينفذ عملية أخذ العينات النموذجية النموذجية للانتشار (DDPM).flaxdiff.samplers.DDIMSampler ): ينفذ عملية أخذ العينات النموذجية الضمنية (DDIM).flaxdiff.samplers.EulerSampler ): عينة محاليل ODE باستخدام طريقة Euler.flaxdiff.samplers.HeunSampler ): عينة حل قصيدة باستخدام طريقة Heun.flaxdiff.samplers.RK4Sampler ): عينة حل قصيدة باستخدام طريقة Runge-Kutta.flaxdiff.samplers.MultiStepDPM ): تنفذ طريقة أخذ العينات متعددة الخطوات مستوحاة من حلال DPM متعدد الخطوات كما هو موضح هنا: Tonyduan/Diffusion) تم تنفيذه في flaxdiff.trainer :
flaxdiff.trainer.DiffusionTrainer ): فئة مصممة لتسهيل تدريب نماذج الانتشار. يدير حلقة التدريب ، وحساب الخسارة ، وتحديثات النموذج. تم تنفيذه في flaxdiff.models :
flaxdiff.models.simple_unet.SimpleUNet ): عينة بنية unet لنماذج الانتشار.flaxdiff.models.simple_unet.Upsample ) ، downsampling ( flaxdiff.models.simple_unet.Downsample ) ، time morseddings ( flaxdiff.models.simple_unet.AttentionBlock ) flaxdiff.models.simple_unet.FouriedEmbedding والكتل المتبقية ( flaxdiff.models.simple_unet.ResidualBlock ). لتثبيت FlaxDiff ، تحتاج إلى الحصول على Python 3.10 أو أعلى. تثبيت التبعيات المطلوبة باستخدام:
pip install -r requirements.txtتم تدريب النماذج واختبارها مع Jax == 0.4.28 والكتان == 0.8.4. ومع ذلك ، عندما قمت بالتحديث إلى أحدث jax == 0.4.30 و Flax == 0.8.5 ، توقفت النماذج عن التدريب. يبدو أن هناك بعض التغييرات الرئيسية في كسر ديناميات التدريب ، وبالتالي أوصي بالالتزام بالإصدارات المذكورة في المتطلبات. txt
فيما يلي مثال مبسط للبدء في تدريب نموذج نشر باستخدام FlaxDiff:
from flaxdiff . schedulers import EDMNoiseScheduler
from flaxdiff . predictors import KarrasPredictionTransform
from flaxdiff . models . simple_unet import SimpleUNet as UNet
from flaxdiff . trainer import DiffusionTrainer
import jax
import optax
from datetime import datetime
BATCH_SIZE = 16
IMAGE_SIZE = 64
# Define noise scheduler
edm_schedule = EDMNoiseScheduler ( 1 , sigma_max = 80 , rho = 7 , sigma_data = 0.5 )
# Define model
unet = UNet ( emb_features = 256 ,
feature_depths = [ 64 , 128 , 256 , 512 ],
attention_configs = [{ "heads" : 4 }, { "heads" : 4 }, { "heads" : 4 }, { "heads" : 4 }, { "heads" : 4 }],
num_res_blocks = 2 ,
num_middle_res_blocks = 1 )
# Load dataset
data , datalen = get_dataset ( "oxford_flowers102" , batch_size = BATCH_SIZE , image_scale = IMAGE_SIZE )
batches = datalen // BATCH_SIZE
# Define optimizer
solver = optax . adam ( 2e-4 )
# Create trainer
trainer = DiffusionTrainer ( unet , optimizer = solver ,
noise_schedule = edm_schedule ,
rngs = jax . random . PRNGKey ( 4 ),
name = "Diffusion_SDE_VE_" + datetime . now (). strftime ( "%Y-%m-%d_%H:%M:%S" ),
model_output_transform = KarrasPredictionTransform ( sigma_data = edm_schedule . sigma_data ))
# Train the model
final_state = trainer . fit ( data , batches , epochs = 2000 )فيما يلي مثال مبسط لإنشاء الصور باستخدام نموذج مدرب:
from flaxdiff . samplers import DiffusionSampler
class EulerSampler ( DiffusionSampler ):
def take_next_step ( self , current_samples , reconstructed_samples , pred_noise , current_step , state , next_step = None ):
current_alpha , current_sigma = self . noise_schedule . get_rates ( current_step )
next_alpha , next_sigma = self . noise_schedule . get_rates ( next_step )
dt = next_sigma - current_sigma
x_0_coeff = ( current_alpha * next_sigma - next_alpha * current_sigma ) / dt
dx = ( current_samples - x_0_coeff * reconstructed_samples ) / current_sigma
next_samples = current_samples + dx * dt
return next_samples , state
# Create sampler
sampler = EulerSampler ( trainer . model , trainer . state . ema_params , edm_schedule , model_output_transform = trainer . model_output_transform )
# Generate images
samples = sampler . generate_images ( num_images = 64 , diffusion_steps = 100 , start_step = 1000 , end_step = 0 )
plotImages ( samples , dpi = 300 ) طراز مدرب على Laion-Aesthetics 12M + CC12M + MS COCO + 1M الجمالية 6+ مجموعة فرعية من Coyo-700m على TPU-V4-32: a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful landscape with a river with mountains, a beautiful forest with a river and sunlight, a beautiful forest with a river and sunlight, a beautiful forest with a river and sunlight, a beautiful forest with a river and sunlight, a beautiful forest with a river and sunlight, a beautiful forest with a river and sunlight, a beautiful forest with a river and sunlight, a beautiful forest with a river and sunlight, a big mansion with a garden, a big mansion with a garden, a big mansion with a garden, a big mansion with a garden, a big mansion with a garden, a big mansion with a garden, a big mansion with a garden, a big mansion with a garden
Params : Dataset: Laion-Aesthetics 12M + CC12M + MS COCO + 1M aesthetic 6+ subset of COYO-700M Batch size: 256 Image Size: 128 Training Epochs: 5 Steps per epoch: 74573 Model Configurations: feature_depths=[128, 256, 512, 1024]
Training Noise Schedule: EDMNoiseScheduler Inference Noise Schedule: KarrasEDMPredictor
الصور التي تم إنشاؤها بواسطة المطالبات التالية باستخدام إرشادات مجانية المصنف مع عامل التوجيه = 2: 'water tulip, a water lily, a water lily, a water lily, a photo of a marigold, a water lily, a water lily, a photo of a lotus, a photo of a lotus, a photo of a lotus, a photo of a rose, a photo of a rose, a photo of a rose, a photo of a rose, a photo of a rose'
Params : Dataset: oxford_flowers102 Batch size: 16 Image Size: 128 Training Epochs: 1000 Steps per epoch: 511
Training Noise Schedule: EDMNoiseScheduler Inference Noise Schedule: KarrasEDMPredictor
الصور التي تم إنشاؤها بواسطة المطالبات التالية باستخدام إرشادات مجانية المصنف مع عامل التوجيه = 4 'water tulip, a water lily, a water lily, a photo of a rose, a photo of a rose, a water lily, a water lily, a photo of a marigold, a photo of a marigold, a photo of a marigold, a water lily, a photo of a sunflower, a photo of a lotus, columbine, columbine, an orchid, an orchid, an orchid, a water lily, a water lily, a water lily, columbine, columbine, a photo of a sunflower, a photo of a sunflower, a photo of a sunflower, a photo of a lotus, a photo of a lotus, a photo of a marigold, a photo of a marigold, a photo of a rose, a photo of a rose, a photo of a rose, orange dahlia, orange dahlia, a lenten rose, a lenten rose, a water lily, a water lily, a water lily, a water lily, an orchid, an orchid, an orchid, hard-leaved pocket orchid, bird of paradise, bird of paradise, a photo of a lovely rose, a photo of a lovely rose, a photo of a globe-flower, a photo of a globe-flower, a photo of a lovely rose, a photo of a lovely rose, a photo of a ruby-lipped cattleya, a photo of a ruby-lipped cattleya, a photo of a lovely rose, a water lily, a osteospermum, a osteospermum, a water lily, a water lily, a water lily, a red rose, a red rose'
Params : Dataset: oxford_flowers102 Batch size: 16 Image Size: 128 Training Epochs: 1000 Steps per epoch: 511
Training Noise Schedule: EDMNoiseScheduler Inference Noise Schedule: KarrasEDMPredictor
Params : Dataset: oxford_flowers102 Batch size: 16 Image Size: 64 Training Epochs: 1000 Steps per epoch: 511
Training Noise Schedule: CosineNoiseSchedule Inference Noise Schedule: CosineNoiseSchedule
Model: UNet(emb_features=256, feature_depths=[64, 128, 256, 512], attention_configs=[{"heads":4}, {"heads":4}, {"heads":4}, {"heads":4}, {"heads":4}], num_res_blocks=2, num_middle_res_blocks=1)
Params : Dataset: oxford_flowers102 Batch size: 16 Image Size: 64 Training Epochs: 1000 Steps per epoch: 511
Training Noise Schedule: CosineNoiseSchedule Inference Noise Schedule: CosineNoiseSchedule
Model: UNet(emb_features=256, feature_depths=[64, 128, 256, 512], attention_configs=[{"heads":4}, {"heads":4}, {"heads":4}, {"heads":4}, {"heads":4}], num_res_blocks=2, num_middle_res_blocks=1)
Params : Dataset: oxford_flowers102 Batch size: 16 Image Size: 64 Training Epochs: 1000 Steps per epoch: 511
Training Noise Schedule: EDMNoiseScheduler Inference Noise Schedule: KarrasEDMPredictor
Model: UNet(emb_features=256, feature_depths=[64, 128, 256, 512], attention_configs=[{"heads":4}, {"heads":4}, {"heads":4}, {"heads":4}, {"heads":4}], num_res_blocks=2, num_middle_res_blocks=1)
لا تتردد في المساهمة من خلال فتح القضايا أو تقديم طلبات السحب. دعنا نجعل Flaxdiff أفضل معا!
هذا المشروع مرخص بموجب ترخيص معهد ماساتشوستس للتكنولوجيا.
