DALLE mtf

DALL-E في شبكة Tensorflow [WIP]

Open-AI's DALL-E في Mesh-Tensorflow.

إذا كان هذا فعالاً بشكل مماثل لـ GPT-Neo، فيجب أن يكون هذا الريبو قادرًا على تدريب نماذج يصل حجمها إلى حجم DALL-E الخاص بـ Open-AI (12B معلمات) وأكبر منه.

لا توجد نماذج مدربة مسبقًا...حتى الآن.

شكرًا لـ Ben Wang على تنفيذ tf vae بالإضافة إلى تشغيل إصدار mtf، وAran Komatsuzaki للمساعدة في بناء mtf VAE وخط أنابيب الإدخال.

يثبت

git clone https://github.com/EleutherAI/GPTNeo
cd GPTNeo
pip3 install -r requirements.txt

إعداد التدريب

يعمل على وحدات TPU، ولم يتم اختباره على وحدات معالجة الرسومات ولكن يجب أن يعمل من الناحية النظرية . تم تصميم تكوينات المثال للتشغيل على جراب TPU v3-32.

لإعداد وحدات TPU، قم بالتسجيل في Google Cloud Platform، وقم بإنشاء مجموعة تخزين.

قم بإنشاء VM الخاص بك من خلال google shell ( https://ssh.cloud.google.com/ ) مع ctpu up --vm-only حتى يتمكن من الاتصال بحاوية Google ووحدات TPU وإعداد الريبو على النحو الوارد أعلاه.

التدريب المسبق على VAE

يحتاج DALLE إلى VAE مدرب مسبقًا لضغط الصور إلى رموز مميزة. لتشغيل التدريب المسبق لـ VAE، اضبط المعلمات في configs/vae_example.json على مسار عالمي يشير إلى مجموعة بيانات من jpgs، واضبط حجم الصورة إلى الحجم المناسب.

  "dataset": {
    "train_path": "gs://neo-datasets/CIFAR-10-images/train/**/*.jpg",
    "eval_path": "gs://neo-datasets/CIFAR-10-images/test/**/*.jpg",
    "image_size": 32
  }

بمجرد الانتهاء من إعداد كل ذلك، قم بإنشاء TPU الخاص بك، ثم قم بتشغيل:

python train_vae_tf.py --tpu your_tpu_name --model vae_example

يسجل التدريب موترات الصورة وقيم الخسارة، للتحقق من التقدم، يمكنك تشغيل:

tensorboard --logdir your_model_dir

إنشاء مجموعة البيانات [DALL-E]

بمجرد تدريب VAE مسبقًا، يمكنك الانتقال إلى DALL-E.

نقوم حاليًا بالتدريب على مجموعة بيانات وهمية. هناك مجموعة بيانات عامة واسعة النطاق لـ DALL-E قيد العمل. في هذه الأثناء، لإنشاء بعض البيانات الوهمية، قم بتشغيل:

python src/data/create_tfrecords.py

يجب أن يقوم هذا بتنزيل CIFAR-10، وإنشاء بعض التسميات التوضيحية العشوائية لتكون بمثابة مدخلات نصية.

يجب تنسيق مجموعات البيانات المخصصة في مجلد، مع ملف jsonl في المجلد الجذر الذي يحتوي على بيانات التسمية التوضيحية والمسارات إلى الصور المعنية، كما يلي:

 Folder structure:

        data_folder
            jsonl_file
            folder_1
                img1
                img2
                ...
            folder_2
                img1
                img2
                ...
            ...

jsonl structure:
    {"image_path": folder_1/img1, "caption": "some words"}
    {"image_path": folder_2/img2, "caption": "more words"}
    ...

يمكنك بعد ذلك استخدام وظيفة create_paired_dataset في src/data/create_tfrecords.py لتشفير مجموعة البيانات في tfrecords لاستخدامها في التدريب.

بمجرد إنشاء مجموعة البيانات، انسخها إلى مجموعتك باستخدام gsutil:

gsutil cp -r DALLE-tfrecords gs://neo-datasets/

وأخيرا، قم بإجراء التدريب مع

python train_dalle.py --tpu your_tpu_name --model dalle_example

دليل التكوين

ضريبة القيمة المضافة:

 {
  "model_type": "vae",
  "dataset": {
    "train_path": "gs://neo-datasets/CIFAR-10-images/train/**/*.jpg", # glob path to training images
    "eval_path": "gs://neo-datasets/CIFAR-10-images/test/**/*.jpg", # glob path to eval images
    "image_size": 32 # size of images (all images will be cropped / padded to this size)
  },
  "train_batch_size": 32, 
  "eval_batch_size": 32,
  "predict_batch_size": 32,
  "steps_per_checkpoint": 1000, # how often to save a checkpoint
  "iterations": 500, # number of batches to infeed to the tpu at a time. Must be < steps_per_checkpoint
  "train_steps": 100000, # total training steps
  "eval_steps": 0, # run evaluation for this many steps every steps_per_checkpoint
  "model_path": "gs://neo-models/vae_test2/", # directory in which to save the model
  "mesh_shape": "data:16,model:2", # mapping of processors to named dimensions - see mesh-tensorflow repo for more info
  "layout": "batch_dim:data", # which named dimensions of the model to split across the mesh - see mesh-tensorflow repo for more info
  "num_tokens": 512, # vocab size
  "dim": 512, 
  "hidden_dim": 64, # size of hidden dim
  "n_channels": 3, # number of input channels
  "bf_16": false, # if true, the model is trained with bfloat16 precision
  "lr": 0.001, # learning rate [by default learning rate starts at this value, then decays to 10% of this value over the course of the training]
  "num_layers": 3, # number of blocks in the encoder / decoder
  "train_gumbel_hard": true, # whether to use hard or soft gumbel_softmax
  "eval_gumbel_hard": true
}

دال-E:

 {
  "model_type": "dalle",
  "dataset": {
    "train_path": "gs://neo-datasets/DALLE-tfrecords/*.tfrecords", # glob path to tfrecords data
    "eval_path": "gs://neo-datasets/DALLE-tfrecords/*.tfrecords",
    "image_size": 32 # size of images (all images will be cropped / padded to this size)
  },
  "train_batch_size": 32, # see above
  "eval_batch_size": 32,
  "predict_batch_size": 32,
  "steps_per_checkpoint": 1000,
  "iterations": 500,
  "train_steps": 100000,
  "predict_steps": 0,
  "eval_steps": 0,
  "n_channels": 3,
  "bf_16": false,
  "lr": 0.001,
  "model_path": "gs://neo-models/dalle_test/",
  "mesh_shape": "data:16,model:2",
  "layout": "batch_dim:data",
  "n_embd": 512, # size of embedding dim
  "text_vocab_size": 50258, # vocabulary size of the text tokenizer
  "image_vocab_size": 512, # vocabulary size of the vae - should equal num_tokens above
  "text_seq_len": 256, # length of text inputs (all inputs longer / shorter will be truncated / padded)
  "n_layers": 6, 
  "n_heads": 4, # number of attention heads. For best performance, n_embd / n_heads should equal 128
  "vae_model": "vae_example" # path to or name of vae model config
}