DALLE mtf
1.0.0
Mesh-Tensorflow の Open-AI の DALL-E。
これが GPT-Neo と同様に効率的である場合、このリポジトリは Open-AI の DALL-E (12B パラメーター) のサイズまで、またはそれ以上のサイズのモデルをトレーニングできるはずです。
事前トレーニングされたモデルはありません...まだです。
tf vae の実装と mtf バージョンの動作に貢献してくれた Ben Wang に感謝します。また、mtf VAE と入力パイプラインの構築を支援してくれた Aran Kokutsuzaki に感謝します。
git clone https://github.com/EleutherAI/GPTNeo
cd GPTNeo
pip3 install -r requirements.txtTPU で実行され、GPU ではテストされていませんが、理論上は動作するはずです。サンプル構成は、TPU v3-32 ポッドで実行するように設計されています。
TPU を設定するには、Google Cloud Platform にサインアップし、ストレージ バケットを作成します。
Google シェル ( https://ssh.cloud.google.com/ ) でctpu up --vm-onlyを使用して VM を作成し、Google バケットと TPU に接続し、上記のようにリポジトリをセットアップできるようにします。
DALLE には、イメージをトークンに圧縮するための事前トレーニングされた VAE が必要です。 VAE 事前トレーニングを実行するには、 configs/vae_example.jsonのパラメータを jpg のデータセットを指す glob パスに調整し、画像サイズを適切なサイズに調整します。
"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_dirVAE が事前トレーニングされたら、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"}
...
次に、 src/data/create_tfrecords.pyのcreate_paired_dataset関数を使用して、トレーニングで使用するためにデータセットを tfrecord にエンコードできます。
データセットが作成されたら、gsutil を使用してバケットにコピーします。
gsutil cp -r DALLE-tfrecords gs://neo-datasets/そして最後に、トレーニングを実行します
python train_dalle.py --tpu your_tpu_name --model dalle_exampleVAE:
{
"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
}
