DALLE mtf

DALL-E en Mesh-Tensorflow [WIP]

DALL-E de Open-AI en Mesh-Tensorflow.

Si esto es igualmente eficiente que GPT-Neo, este repositorio debería poder entrenar modelos hasta y más grandes que el tamaño de DALL-E de Open-AI (12B parámetros).

No hay modelos previamente entrenados... todavía.

Gracias a Ben Wang por la implementación de tf vae y por hacer funcionar la versión mtf, y a Aran Komatsuzaki por su ayuda en la construcción de mtf VAE y el canal de entrada.

Configuración

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

Configuración de entrenamiento

Se ejecuta en TPU, no se ha probado en GPU, pero debería funcionar en teoría . Las configuraciones de ejemplo están diseñadas para ejecutarse en un pod TPU v3-32.

Para configurar TPU, regístrese en Google Cloud Platform y cree un depósito de almacenamiento.

Cree su VM a través de un shell de Google ( https://ssh.cloud.google.com/ ) con ctpu up --vm-only para que pueda conectarse a su depósito de Google y a sus TPU y configurar el repositorio como se indica arriba.

Preentrenamiento VAE

DALLE necesita un VAE previamente entrenado para comprimir imágenes en tokens. Para ejecutar el preentrenamiento de VAE, ajuste los parámetros en configs/vae_example.json a una ruta global que apunte a un conjunto de datos de jpg y ajuste el tamaño de la imagen al tamaño apropiado.

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

Una vez que todo esto esté configurado, cree su TPU y luego ejecute:

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

El entrenamiento registra tensores de imagen y valores de pérdida; para verificar el progreso, puede ejecutar:

tensorboard --logdir your_model_dir

Creación de conjuntos de datos [DALL-E]

Una vez que el VAE esté previamente entrenado, puede pasar a DALL-E.

Actualmente estamos entrenando con un conjunto de datos ficticio. Se está trabajando en un conjunto de datos públicos a gran escala para DALL-E. Mientras tanto, para generar algunos datos ficticios, ejecute:

python src/data/create_tfrecords.py

Esto debería descargar CIFAR-10 y generar algunos subtítulos aleatorios para que actúen como entradas de texto.

Los conjuntos de datos personalizados deben formatearse en una carpeta, con un archivo jsonl en la carpeta raíz que contenga datos de título y rutas a las imágenes respectivas, de la siguiente manera:

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

Luego puede usar la función create_paired_dataset en src/data/create_tfrecords.py para codificar el conjunto de datos en tfrecords para usarlo en el entrenamiento.

Una vez creado el conjunto de datos, cópialo en tu depósito con gsutil:

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

Y finalmente, ejecuta el entrenamiento con

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

Guía de configuración

VAE:

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

DALL-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
}