DALLE mtf

DALL-E ใน Mesh-Tensorflow [WIP]

DALL-E ของ Open-AI ใน Mesh-Tensorflow

หากสิ่งนี้มีประสิทธิภาพใกล้เคียงกับ GPT-Neo repo นี้ควรจะสามารถฝึกโมเดลได้มากถึงและใหญ่กว่าขนาดของ DALL-E (พารามิเตอร์ 12B) ของ Open-AI

ยังไม่มีโมเดลที่ผ่านการฝึกอบรม...

ขอขอบคุณ Ben Wang สำหรับการใช้งาน tf vae ตลอดจนการทำให้เวอร์ชัน mtf ทำงาน และ Aran Komatsuzaki สำหรับความช่วยเหลือในการสร้าง mtf VAE และไปป์ไลน์อินพุต

ตั้งค่า

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

การตั้งค่าการฝึกอบรม

ทำงานบน TPU ซึ่งยังไม่ทดสอบบน GPU แต่ควรใช้งานได้ ในทางทฤษฎี การกำหนดค่าตัวอย่างได้รับการออกแบบมาให้ทำงานบนพ็อด TPU v3-32

หากต้องการตั้งค่า TPU ให้ลงชื่อสมัครใช้ Google Cloud Platform และสร้างที่เก็บข้อมูล

สร้าง VM ของคุณผ่าน Google Shell ( https://ssh.cloud.google.com/ ) ด้วย ctpu up --vm-only เพื่อให้สามารถเชื่อมต่อกับที่เก็บข้อมูล Google และ TPU ของคุณและตั้งค่า repo ข้างต้น

การฝึกอบรม VAE ล่วงหน้า

DALLE ต้องการ VAE ที่ได้รับการฝึกล่วงหน้าเพื่อบีบอัดรูปภาพเป็นโทเค็น หากต้องการรันการฝึกล่วงหน้าของ VAE ให้ปรับพารามิเตอร์ใน configs/vae_example.json เป็นพาธ glob ที่ชี้ไปยังชุดข้อมูล 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
}

ดอล-อี:

 {
  "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
}