More "OpenAI Blog Post" Training | Depth 32 | Heads 8 | LR 5e-4

[afiaka87]

Edit: Moved to discussions: #106

Hey, all. Some of you might know I'm practicing and learning about machine learning with dalle-pytorch and a dataset consisting of the images OpenAI presented in the DALLE blog post. I honestly dont have the money to train this whole dataset,

edit: this is no longer true. Using the 1024 VQGAN from the "Taming Transformers" research, it's now quite possible to train a full dataset of 1,000,000 image-text pairs and i'm doing just that. I hope to have it finished in about a week. I assume someone else will release a dalle-pytorch trained properly on COCO and other image sets before then, but if they dont, check here for updates.

Anway, it ran for ~36000 steps. As you can see it...still really likes mannequins. I'm considering removing them from the dataset. But also, you'll notice that the network has actually got a decent idea of the sort of general colors that belong in types of prompts.

Some Samples from Near the End of Training

Every Text-Image Reconstruction

https://wandb.ai/afiaka87/dalle_pytorch_live_training/reports/dalle-pytorch-Test-Run-2--Vmlldzo1MzM5MjQ

Deliverables (my train_dalle.py)

https://gist.github.com/afiaka87/850fb3cc48edde8a7ed4cb1ce53b6bd2

This has some code in it that actually manages to deal with truncated images via Try Catch. Apparently detecting a corrupted PNG is harder than P vs NP. PIL's imverify() function doesnt catch all of them. Python's built in imghdr library doesn't catch all of them either. So you just sort of catch OSError and return an item further along. Works well enough.

Parameters

SHUFFLE = True
EPOCHS = 28 # This wound up being less than a single epoch, of course. 
BATCH_SIZE = 16
LEARNING_RATE = 0.0005 # I found this learning rate to be more suitable than 0.0003 in my hyperparameter sweep post
GRAD_CLIP_NORM = 0.5
DEPTH = 32
HEADS = 8
MODEL_DIM = 512
TEXT_SEQ_LEN = 256
DIM_HEAD = 64
REVERSIBLE = True,
ATTN_TYPES = ('full')

Dataset Description

#61 (comment)

Just for more info on the dataset itself, it is roughly 1,100,000 256x256 image-text pairs that were generated by OpenAI's DALL-E. They presented roughly ~30k unique text prompts of which they posted the top 32 of 512 generations on https://openai.com/blog/dall-e/. Many images were corrupt, and not every prompt has a full 32 examples, but the total number of images winds up being about 1.1 million. If you look at many of the examples on that page, you'll see that DALL-E (in that form at least), can and will make mistakes. These mistakes are also in this dataset. Anyway I'm just messing around having fun training and what not. This is definitely not going to produce a good model or anything.

There are also a large number of images in the dataset which are intended to be used with the "mask" feature. I don't know if that's possible yet in DALLE-pytorch though. Anyway, that can't be helping much.

[afiaka87]

@lucidrains btw, i've been going through the wandb.ai docs and found some nice extras you can add to train_dalle.py that will give you live updates on the transformer itself:

config = wandb.config
config.depth = DEPTH
config.heads = HEADS
config.dim_head = DIM_HEAD
config.learning_rate = LEARNING_RATE
config.shuffle = SHUFFLE
config.resume = RESUME
config.batch_size = BATCH_SIZE
config.grad_clip_norm = GRAD_CLIP_NORM
config.reversible = REVERSIBLE
config.model_dim = MODEL_DIM
config.attn_types = ATTN_TYPES

wandb.init(project = PROJECT_NAME, resume = RESUME)

wandb.watch(dalle) # Updates a graph of gradients on wandb as soon as your model changes

In particular, that very last line is actually all you need to add. But attaching all the parameters in the way I did also allows it to track those better and you can more easily create hyperparameter sweeps from existing projects when you do.

[lucidrains]

@afiaka87 ohh got it! i'm circling back to DALL-E this week for some final instrumentations :) i'll be sure to add that! 🙏

[afiaka87]

@lucidrains Awesome, looking forward to it! Thanks for patching up big-sleep/deep-daze btw. I tried but I'm so distracted with this project now lol.

[lucidrains]

@afiaka87 yes, arguably getting a DALL-E model trained and released would be bigger than either big sleep or deep daze!

[lucidrains]

thanks for doing this! it demonstrates that reversibility does work :)

[afiaka87]

@afiaka87 yes, arguably getting a DALL-E model trained and released would be bigger than either big sleep or deep daze!

@lucidrains For sure! I'm trying to be as open about my training, code, results, etc. but I'm not seeing much else of that here. I'm aware it's prohibitively expensive for most though and I'm privileged to be able to run Depth=32 for a day or two. At any rate, looking forward to 1024 token model from Germany! I know it's in there currently but I'm still having some trouble with it last I checked. All in due time.

[afiaka87]

thanks for doing this! it demonstrates that reversibility does work :)

There should be system usage info in the graphs on wandb.ai, but yeah it does what it says on the label, lol. You definitely trade time for space. But, that whole training session never went above 16 GiB of VRAM. So at least people can use colab!

[lucidrains]

@afiaka87 great to know! also, do log the issue with the VQ-GAN VAE and i'll be sure to fix it this week. It seems to be working on my end, but I haven't tried testing it from a fresh install

[afiaka87]

@lucidrains one last thing, but the "image masking" feature is used pretty thoroughly in this dataset and they even have the image used for the mask and everything. Let me know as soon as that feature is implemented as I would love to use those as a baseline for that.