Repository navigation
Dreambooth broken, possibly because of ADAM optimizer, possibly more. #712
Description
Activity
Hi @affableroots , thanks for the issue!
What do you mean by the whole model still degrades ?
The dreambooth paper uses imagen model, which keeps the text encoder frozen. The VAE should also be frozen.
As far as I can see this https://github.com/JoePenna/Dreambooth-Stable-Diffusion also uses prior preservation. Did you try training with this and did it work ? I would be happy to investigate.
Also note that, for dreambooth we need to play around a bit with hyperparamters i.e the number of steps, learning rate, number of prior images etc. Could you please post more details
- How many images are you using for training
- How many priro images
- The number of training steps
- The instance prompt and class prompt
- And if possible the training images that you are using.
This will help us debug the issue better. Thank you.
If I understand correctly, this paper says they finetune the entire model, although in my experiments I agree and have kept the VAE frozen. I guess the following could be saying they only finetune all layers down stream of the text embeddings?
RE JoePenna, I think his repo is an example of Textual Inversion + unfrozen unet. Quoted from his readme:
We're [now] realizing the "regularization class" bits of this code have no effect, and that there is little to no prior preservation loss.
So, out of respect to both the MIT team and the Google researchers, I'm renaming this fork to: "Unfrozen Model Textual Inversion for Stable Diffusion".
For an alternate implementation that attempts prior loss preservation, please see "Alternate Option" below.
This repo currently gives far better results than the huggingface version, but I think the HF version has more promise once fixed.
RE my params, I've tried ~50ish permutations within the following bounds, but in general I've tried to stick close to the example params in your readme, and the params in the DB paper.
- Batches: 1-2
- 8bit optimizer on/off
- Number of images: 5 - 50
- Number of prior images: 200-1500
- Number of steps: 500 - 5000
- Learning rate: 1e-7 to 1e-4
- Prior Preservation Loss on/off
- Prompts are always of the form: "a photo of CLASS" and "a photo of TOKEN CLASS", where TOKEN is
sks(struggles to train away from rifles), or random three letter words, or, more recently, I added textual inversion to your example to add a token to the tokenizer. - I've tried many types of images, but have settled on a test case of simply learning faces that aren't already in SD, eg, my face.
RE degradation, I mean the whole model slowly becomes unusable over the course of learning. For instance given the stock params of 500 steps * lr=5e-6, it poorly learns the instance face, and all the class images visually degrade somewhat. It forgets things it knows already, and generally turns to visual noise, at any learning rate, and for any number of steps, in or out of the finetuning set. Very low learning rates, eg 1e-7, slow the process, but then you're not learning your instance either.
For inference I always test with: "a photo of token class", " a black and white sketch of token class", "a photo of johnny depp". So I've got in there: the exact prompt I trained with, a prompt that tests generalization, and a non-instance prompt to see if it's forgetting other faces. Johnny Depp degrades over the course of learning.
Given my limited understanding, could this all be due to weight regularization performed by Adam, decaying the weights? I seem to have better results using SGD, although it learns slowly.
I'm still getting quite poor results training with this library's (huggingface's) DB example.
I've been using ShivamShrirao's fork of this repo for several days now. With my own Google Colab notebook (which adds additional configuration options & streamlines the whole process), I've managed to generate some pretty decent renders of my younger self...
Some of the renders generated
Actual photos of me used as input for the training process
@jslegers Thanks for the follow up! I've had results look fairly similar, but I think everythings still degrading at least on my end. For instance:
- "a photo of" type prompts look a bit overcooked for the trained instance, as well as out-of class instances (unless my lr is low enough to not learn the instance)
- How do results look for random actors, compared to pre-training? Mine tend to degrade considerably, and often look "overcooked", ie that look things get at real high CFG.
- And how about using your likeness in non-portrait situations like, playing a sport, or an instrument, etc.?
Oh, another symptom that I think indicates something is wrong, is I have not been able to overtrain. What I mean is, given a very high lr and/or high step count, everything sort of degrades without looking like my new instance. Faces look melted, and drift away from their priors, eg actors faces are no longer their faces, and they probably "melt" and become very indistinct, as faces. I notice photos degrade quicker than artwork.
Contrarily, I feel like JoePenna's version is easy to overtrain. Sure the trained concept leaks into everything, but I'm using this as a sort of diagnostic tool. In his version, an overtrained concept looks unmistakably like that concept (even if it increasingly fails to generalize)
Is there no merit to the weight decay hypothesis? That still sticks out in my mind. The way I see it is if you have a whole universe of latent space, but you only update the small corner that looks like your face, but regardless you regularize the entire thing towards 0, you're going to degrade everything, which would explain why overtraining doesn't work, generalization, learning a concept in the first place, etc.
Given my hypothesis, with a low step count you could learn your concept without degrading things too much, as @jslegers and others (and myself) have demonstrated, but it still doesn't seem like a quality finetuning for all the reasons I've mentioned, IE how the model degrades.
Or I could be crazy and running this all wrong?
but I think everythings still degrading at least on my end.
For my use case, I noticed that the diversity with the class "man" started declining significiantly. At some point, all random class images generated during the training process look identical.
I haven't really looked much into the degrading of my "environment", though, as I was too focused to trying to fix my likeness. But it does indeed appear to be degrading as training increases.
How do results look for random actors, compared to pre-training? Mine tend to degrade considerably, and often look "overcooked", ie that look things get at real high CFG.
Is this that "overcooked" look you're talking about?
Nightmare fuel, if you ask me...
And how about using your likeness in non-portrait situations like, playing a sport, or an instrument, etc.?
I haven't really invested much time in doing full body renders, but my best models did produce some OK results :
he way I see it is if you have a whole universe of latent space, but you only update the small corner that looks like your face, but regardless you regularize the entire thing towards 0
That seems to be what's happening, yes.
As a sidenote, I'm struggling to find a way to build upon previous trainings. When I have a general "johnslegers" concept, for example, I'd like to build upon that training be training eg. the concept "close-up or johnslegers" or "detailed portrait of johnslegers", as to refine the model further. But this doesn't seem possible. Attempts to achieve this by tinkering with my class prompt, instance prompt, class name or concept name generally leads to a loss of previous training achieved and serious degrading the quality of my models.
Thank you for running that again! It seems we have similar results, though, your generalization attempts do look better than what I've been achieving, even given the stock params from the readme. And yes, that's the overcooking I'm talking about. That seems specifically like certain latent values are... clipping, if that's the right term.
try to refine the model further
Ya imagine that's impossible if the entire latent space is getting regularized and forgetting everything it learned in its original training. I've turned off
weight_decayand while it seems to help, the whole latent space gets degraded anyway.So then I thought maybe the Prior Preservation Loss was affecting things incorrectly, and I turned that off. So, ostensibly it should become a run of the mill finetuning of just unet, I think. I would expect that training without PPL would simply, for instance, pull all faces toward your face, but not destroy everything. But, all faces somehow just get cooked!
Hypothesis: I'm strongly suspecting weight decay is an issue and some other hidden factor.
Ya imagine that's impossible if the entire latent space is getting regularized and forgetting everything it learned in its original training.
Probably.
I guess I better just start all over with the original SD model and play around with different image files & config settings to reproduce at minimum the quality I achieved to far than continue tinkering with existing models...
@affableroots In the paper by all layers they mean all layers of the unet model trained, the rest i.e the text encode is kept frozen.
Looking at the discussion and also from my experiments the model does run into overfitting and catastrophic forgetting. To me this seems likely as we finetune the whole model with few images. For now I don't know any way to avoid this, I am going to discuss this with the author and let you know if we find anything.
To sum it up, I'm not sure if the script is broken, it's rather that the model quickly overfits.
I appreciate your attention on this @patil-suraj!
Looking at the discussion and also from my experiments the model does run into overfitting and catastrophic forgetting. To me this seems likely as we finetune the whole model with few images. For now I don't know any way to avoid this, I am going to discuss this with the author and let you know if we find anything.
I suspect overfitting issues and other related issues can be reduced by finetuning your number of training steps, class prompt, number of variety of class images generated, quality and variety of your input images, etc.
I've been experimenting with various strategies : use just a handful of pics, use many pics, auto-generate the class pics, choose class pics myself, do only 250 training steps per run, do 2000 training steps each time, etc.
My results after just a single run from the original SD model vary significantly depending on various of these criteria and I have yet to determine any pattern.
We should develop some best practices on how to optimize the model we're creating while keeping degradation of the latent space minimal. If anyone should be able to provide some info we all seem to be missing, that's going to be the original author...
While we're exploring the tuning of DB from all angles, I'll mention some recently decent results. I'm using my modified Dreambooth + Textual Inversion:
- 5 new tokens+embeddings
- 2000 steps
- lr5e-5 for text embeddings
- lr5e-6 for unet
- weight decay=0.0
- prior loss weight=0.1
- no. class images=1500
I will say that faces are still getting cooked, so, I'm going to play with this some more.
My theory is that if you combine DB + TI, you can ask TI to learn more strongly without risking damage to the rest of the model (because you're only learning brand new embeddings that don't cause other domains to forget), but still train unet to help out a little, where TI is struggling.
I'm using my modified Dreambooth + Textual Inversion
Any chance you can make your modified script available for the rest of us, so we can test it ourselves?
13 remaining items
Been testing using my face all day with a variety of optimisers because reasons :) Things learnt so far -
- The defaults seem very good so far, with 800-1800 steps being OK for a face. Gives perfectly useable results in 10 minutes.
- However, my face bleeds into all the other faces. Some (Abraham Lincoln, Donald Trump) are impacted less
- It for sure changes the whole model regardless of "class". This can be easily seen by switching to something like DiffGrad.
- Slightly lower LR (4e-6) with more steps (1600) and more preservation images seems to lessen the leakage into other faces (even if the class is just a symbol):
Default Johnny Depp (seed 1201043):
With DB model (defaults), 200 preservation images, fp16, 1000 steps:
I can really see my features in there!With DB model, 1200 preservation images, fp16, 1600 steps, 4e-6:
Not so much me :)Reacted by Wes Roberts@nerdyrodent, have you tried 1e-6 for 2000 steps which is basically the JoePenna's default training parameters.
@compustar - just tried 2000 steps at 1e-6 and the results weren't as good as the default 800 steps at 5e-6.
Great thread you started here @affableroots
I've been trying to run your TI-Dreambooth-Tokenizer code. I am getting errors about accessing huggingface-cli that I am not getting when I run Shivam's version of the code. specifically:
OSError: CompVis/stable-diffusion-v1-4 is not a local folder and is not a valid model identifier listed on 'https://huggingface.co/models'
If this is a private repository, make sure to pass a token having permission to this repo withuse_auth_tokenor log in withhuggingface-cli login.
Traceback (most recent call last):I am logged into hugface-cli tried; tried different tokens with no success.
I tried downloading the files locally and pathing them there, but eventually it asks for a config file i cannot find, however this seems to be an access problem rather than a file location problem.do you (or anyone) have an updated version, or do you know how i can fix this?
JoePenna "Dreambooth"
You need to download the diffusers model, which is a repo with multiple directories and files. The .ckpt is for regular stable diffusion, etc.
Think @patil-suraj will very soon release a nice update of dreambooth :-)
Reacted by affableroots, InB4DevOps, bjoerm and David Alejandro@patil-suraj Did you figure out what the issue was then? I've been curious.
It might be related to this PR: #883
Hi everyone! Sorry to be so late here.
We ran a lot of experiments with the script to see if there's any issues or if it's broken.
Turns out, we need to carefully pick hyperparameters like LR and training steps to get good results with dreambooth.Also, the main reason the results with this script were not as good as Compvis forks, is that the
text_encoderis being trained in those forks, and it makes a big difference in quality especially on faces.We compiled all our experiments in this report, and also added an option to train the
text_encoderin the script which can be enabled by passing the--train_text_encoderargument.Note that, if we train the
text_encoderthe training won't fit on 16GB GPU anymore, it will need at least 24GB VRAM. It should still be possible to do it 16GB using deepspeed, but will be slower.Please take a look at the report, hope you find it useful.
Reacted by Patrick von Platen, Michał Zając, Wes Roberts, ooofest and Adam ŽingorI tested ShivamShirao's fork with a 12GB card and it trained without errors, It seems that the text_encoder's weight have been updated. I haven't had time to compare the results of the two forks.
My parameters: with 8bit_adam, without caching latents, with prior preservation, without mixed precision, with gradient checkpointing, with text encoder training.- I'll see if I can implement that one today.…On Wed, Oct 19, 2022, 5:38 AM matteo ***@***.***> wrote: I tested ShivamShirao's fork with a 12GB card and it trained without errors, I seems that the text_encoder's weight have been updated. I haven't had time to compare the results of the two forks. My parameters: with 8bit_adam, without caching latents, with prior preservation, without mixed precision, with gradient checkpointing, with text encoder training. — Reply to this email directly, view it on GitHub <#712 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAMO4NBB46SJFKYV7ANGJOTWD7FTRANCNFSM6AAAAAAQ4C6WVE> . You are receiving this because you commented.Message ID: ***@***.***>
Thanks for the review @patil-suraj! It sounds like the difference was all in unfreezing the text embeddings. It sounds like lr+nsteps makes a small target to hit, and your experiment helps us aim better.
@affableroots I am getting good outputs for Dreambooth by training the images of a certain face. However I am not getting consistently good results i.e the model doesn't product the similar face. It changes it a little bit in some instances.
Even when I take the best output and just change the prompt, it changes the face a bit.Any idea on how to solve the issue?
I think Huggingface's Dreambooth is the only popular SD implementation that also uses Prior Preservation Loss, so I've been motivated to get it working, but the results have been terrible, and the entire model degrades, regardless of: # timesteps, learning rate, PPL turned on/off, # instance samples, # class regularization samples, etc. I've read the paper, and found that they actually unfreeze everything including the text embedder (and VAE? I'm not sure so I leave it frozen), so I implemented textual inversion within the dreambooth example (new token, unfreeze a single row of the embedder), which improves results considerably, but the whole model still degrades no matter what.
Someone smarter than me can confirm, but I think the culprit is ADAM:
diffusers/examples/dreambooth/train_dreambooth.py
Line 426 in f1484b8
My hypothesis is that since ADAM tries to drag all weights of unet etc. to 0, it ruins parts of the model that aren't concurrently being trained during the finetuning.
I've tested with
weight_decayset to 0, and results seem considerably better, but I think the entire model is still degrading. I'm trying SGD next, so, fingers crossed, but there may still be some dragon lurking in the depths even despite removing ADAM.A point of reference on this journey is the JoePenna "Dreambooth" library which doesn't implement PPL, and yet preserves priors much much better than this example, not to mention it learns the instance better, and is far more editable, and preserves out-of-class rather well. I expect more from this huggingface dreambooth example, and I'm trying to find why it's not delivering.
Any thoughts or guidance?
EDIT1A: SGD didn't learn the instance at 1000 steps + lr=5e-5, but it definitely preserved the priors way better (upon visual inspection. The loss really doesn't decrease much in any of my inversion/dreambooth experiments).
EDIT1B: Another test failed to learn using SGD at 1500 steps + lr=1e-3 + momentum=0.9. It might be trying to learn, but, not much. Priors were nicely preserved though still.
EDIT1C: 1500 * lr=5e2 learned great, was editable, didn't destroy other priors!!!
EDIT2: JoePenna seems to use AdamW, so I'm not sure what's up anymore, but I'm still getting quite poor results training with this library's (huggingface's) DB example.