DALLE mtf

DALL-E dans Mesh-Tensorflow [WIP]

DALL-E d'Open-AI dans Mesh-Tensorflow.

Si cela est tout aussi efficace que GPT-Neo, ce dépôt devrait être capable de former des modèles jusqu'à et plus grands que la taille du DALL-E d'Open-AI (paramètres 12B).

Pas de modèles pré-entraînés... Pour l'instant.

Merci à Ben Wang pour l'implémentation de tf vae ainsi que pour le fonctionnement de la version mtf, et à Aran Komatsuzaki pour son aide à la construction du mtf VAE et du pipeline d'entrée.

Installation

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

Configuration de la formation

Fonctionne sur des TPU, non testé sur des GPU mais devrait fonctionner en théorie . Les exemples de configuration sont conçus pour s'exécuter sur un pod TPU v3-32.

Pour configurer des TPU, inscrivez-vous à Google Cloud Platform et créez un bucket de stockage.

Créez votre VM via un shell Google ( https://ssh.cloud.google.com/ ) avec ctpu up --vm-only afin qu'elle puisse se connecter à votre bucket Google et à vos TPU et configurer le dépôt comme ci-dessus.

Préformation VAE

DALLE a besoin d'un VAE pré-entraîné pour compresser les images en jetons. Pour exécuter le pré-entraînement VAE, ajustez les paramètres dans configs/vae_example.json sur un chemin global pointant vers un ensemble de données de jpg et ajustez la taille de l'image à la taille appropriée.

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

Une fois tout cela configuré, créez votre TPU, puis exécutez :

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

La formation enregistre les tenseurs d'image et les valeurs de perte, pour vérifier la progression, vous pouvez exécuter :

tensorboard --logdir your_model_dir

Création d'un ensemble de données [DALL-E]

Une fois le VAE préformé, vous pouvez passer au DALL-E.

Actuellement, nous nous entraînons sur un ensemble de données factices. Un ensemble de données publiques à grande échelle pour DALL-E est en préparation. En attendant, pour générer des données factices, exécutez :

python src/data/create_tfrecords.py

Cela devrait télécharger CIFAR-10 et générer des légendes aléatoires pour servir de saisie de texte.

Les ensembles de données personnalisés doivent être formatés dans un dossier, avec un fichier jsonl dans le dossier racine contenant les données de légende et les chemins d'accès aux images respectives, comme suit :

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

vous pouvez ensuite utiliser la fonction create_paired_dataset dans src/data/create_tfrecords.py pour encoder l'ensemble de données dans tfrecords pour une utilisation dans la formation.

Une fois l'ensemble de données créé, copiez-le dans votre bucket avec gsutil :

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

Et enfin, organisez une formation avec

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

Guide de configuration

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
}