20 Epochs on COCO - (Larger Transformer)

[afiaka87]

This one was done entirely on one GPU over the course of about 48 hours. It has no dependencies on DeepSpeed and was trained with 32 bit precision.

If you want to train your own DALL-E and have about 8 GiB of VRAM; this is a very good configuration in my opinion.

Full W&B training session

Run inference with the following notebook:
https://colab.research.google.com/drive/1UqTIilYvILiO_mH88LP67YkBRmc0dEFR

Details
Transformer:

• Visual Dim - 512
• Max Text Length/Language Dim - 80
• Depth - 16
• Heads - 16
• Head Dim - 64
• Attention: (axial_row, axial_row, axial_col, axial_row,axial_row, axial_row, axial_col, axial_row,axial_row, axial_row, axial_col, full,axial_row, axial_row, axial_col, full) # note the two layers of dense attention
• Reversible = False

VQGAN:

• Pretrained ImageNet 1024 (f=16)

Training

• lr_decay = True
• learning_rate 3e-4
• batch size = 9 # just barely fits.
• ga_steps - 1

Augmentations:

• RandomResizeRecrop
  • lower bound - 1.0
  • interpolation_mode = BILINEAR
• ColorJitter
  • brightness=0.025
  • hue=0.0
  • contrast=0.025
  • saturation=0.2
  • Slight variation in contrast/brightness with a focus on a change in saturation; which I assume to continue to be an accurate representation of the captions with this particular dataset
• RandomHorizontalFlip
  • 10 percent chance of hflip.
  • Messes with text but this dataset doesnt have much and many captions still make sense when flipped (unlike, for instance, Open Images Localized Annotations - which makes heavy usages of the words left,right,top,bottom to describe where things are on the image.

Hardware:

• 1x RTX 2070 Super (8 GiB VRAM)
• 1x AMD Ryzen 3900
• 32 GiB RAM

Checkpoints:
You can find a checkpoint for each epoch trained here:
https://wandb.ai/dalle-pytorch-replicate/COCO512_16_16D_16H_80TSL/artifacts/model/trained-dalle/07c445559fd9183e302e

[rom1504]

I'm curious of the impact to use these axial row and axial col layers vs full attn. I might try it too

[afiaka87]

I actually tried using only full attention on this one because I had indeed seen it help with the loss a little with the GumbelVQ.

Didn't fit into VRAM. This happened to be the very next configuration I tried which did. So at least for a single card; it's been useful.

[rom1504]

nice results!

[afiaka87]

Thanks! For the dataset size it's certainly a bit surprising. But it definitely struggles with out-of-corpus captions. There's a grand canyon example in the colab notebook which works rather well on the CC12M checkpoint but simply fails on this one. I don't think there's any shots of it in this dataset?

[wintersurvival]

Nice results!
The inference parameter "--bpe_path variety.bpe" seems that you use a different vocab, could you tell me what you use?

[afiaka87]

wget https://www.dropbox.com/s/a95u0mw5bwp2hr7/new_variety.bpe?dl=1 -O variety.bpe

[afiaka87]

@wintersurvival

[wintersurvival]

Thank you @afiaka87

[wintersurvival]

The vocab size of the variety.bpe using YttmTokenizer is 8192, amazing!
Did you train the bpe using captions of MS COCO dataset?

[afiaka87]

@wintersurvival Very sorry for the delay on the response. If I recall the tokenizer was trained on about half of the captions used in CC12M (roughly ~6 million). I also included a lot of captions from crawling@home, unfiltered, to give the tokenizer some notion of "internet-ish" text downstream.

Doing so is trivial with youtokentome. The repo also supports huggingfaces which of course has a wide variety of fancy tokenizers as well.

pip install youtokentome

gives you access to the cli tool yttm. Your captions need to be in a single line separated txt file.

 yttm bpe --help
Usage: yttm bpe [OPTIONS]

  Train BPE model.

Options:
  --data PATH           Training data file path.  [required]
  --model PATH          Output model file path.  [required]
  --vocab_size INTEGER  Number of tokens in the final vocabulary.  [required]
  --coverage FLOAT      Percentage of characters covered by the model.
                        [default: 1.0]
  --n_threads INTEGER   Number of threads.  [default: -1]
  --pad_id INTEGER      Padding token id.  [default: 0]
  --unk_id INTEGER      Unknown token id.  [default: 1]
  --bos_id INTEGER      Begin of sentence token id.  [default: 2]
  --eos_id INTEGER      End of sentence token id.  [default: 3]
  --help                Show this message and exit.

[wintersurvival]

[afiaka87]

That sounds like something must have changed in the codebase since i did this run.

[Asthestarsfalll]

@afiaka87
Hi, I download the checkpoint from the link you give, but it shows that the dalle only has a 8-layers transformer with attention type 'sparse