[Diffusers] [SDv1.5 - Text-to-Image Finetuning] Choose the best checkpoint for inference.

[VietAnhNguyen20]

Search before asking:

• I have thoroughly reviewed the Diffusers and Huggingface documentation, the open or closed issues on github
  , as well as examined the code provided by Huggingface. Additionally, I have explored related resources concerning the process of selecting the optimal checkpoint (UNET checkpoint) for Stable Diffusion after finetuning:
  • Generating images for each epoch using the validation prompt (a qualitative method, derived from the code provided by Huggingface).
  • Assessing all checkpoints using the FID or CLIP metrics (a quantitative method, which I found suitable, as outlined in the Huggingface documentation regarding the validation process).
• However, I have encountered no related issues, fragments of code, or documentation regarding validation loss during training.

Context:

Version of packages used:

• I am utilizing diffusers (0.27.0.dev0) and transformers (4.38.1).

Purpose of training:

• I am working with a substantial dataset (~600,000 pairs of (text, image) for training), hence I have opted to finetune the entire model (SDv1.5) at a lower learning rate, rather than employing LoRA. Additionally, I possess a validation dataset (~60,000 pairs of (text, image)). To accommodate my data, I have already developed a custom Dataset.

Command for training:

accelerate launch --mixed_precision="fp16" train_text_to_image.py --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" --use_ema --resolution=224 --center_crop --random_flip --train_batch_size=128 --gradient_accumulation_steps=1 --gradient_checkpointing --num_train_epochs=10 --learning_rate=1e-05 --max_grad_norm=1 --lr_scheduler="constant" --lr_warmup_steps=0 --output_dir="Testing" --image_column="image_path" --enable_xformers_memory_efficient_attention --snr_gamma=5 --validation_epochs=1 --validation_prompts="Some qualitative prompts to test"

Results:

• The model demonstrates the capability to learn the "style" corresponding to my dataset, evident from its performance from the initial epoch onwards (I trained for 10 epochs).

Question:

• After training my model for approximately 10 epochs and logging a checkpoint every 1500 steps, I am uncertain about the best checkpoint for inference. Typically, I would consider both the train loss and the validation loss to select a few optimal checkpoints for further analysis. However, the code for diffusers does not include validation loss, and I am concerned that selecting the final checkpoint might lead to overfitting issues.

• Qualitative method:

  • I have implemented a validation prompt to generate images at each epoch, enabling qualitative assessment. However, given that the model exhibits substantial learning from the initial epoch, I find that all checkpoints produce "reasonable" outputs for certain categories of input texts. Nonetheless, due to the complexity of my problem, I aim to identify the best model or the most favorable trade-off checkpoint for inference.

• Quantitative method:

  • CLIP evaluation between original images and images generated from corresponding texts (comparing image embeddings):

    • I have computed similarities for each pair of original/generated images and constructed an estimated Gaussian distribution. While observing the Mean and Std, it's apparent that the model has learned, as indicated by the shifting curve. However, due to the marginal difference in values, pinpointing the optimal checkpoint might be challenging. I have provided the corresponding plot for reference.
    • Notably, the checkpoint yielding the best result in terms of Mean does not necessarily coincide with the one identified using the FID metric.

  • CLIP evaluation between generated images and corresponding texts (comparing image and textual embeddings):

    • Given that my textual inputs are not highly representative and resemble code snippets, this metric may not be optimal for identifying the best checkpoint. Nevertheless, I have included the plot for discussion purposes.

  • FID evaluation between a generated collection of images based on texts and a collection of original images (FID calculation using torchmetrics):

    • Utilizing the FID metric outlined in the Huggingface tutorial, which employs FID calculation from torchmetrics, I have observed a relatively high score (~80) due to the fact that InceptionNet v3 is trained on ImageNet features. While this metric offers a potential means of checkpoint selection, I am apprehensive about its suitability for my specific problem. My dataset predominantly consists of "anime style" data, and I aim to generate stylized rather than realistic images.

• In conclusion, the absence of validation loss implementation or discussion prompts me to seek guidance from our community on the optimal approach for selecting checkpoints. Despite conducting thorough research on this matter, my knowledge remains limited, and I am eager to explore additional perspectives and insights on this issue.

Related figures:

[sayakpaul]

After training my model for approximately 10 epochs and logging a checkpoint every 1500 steps, I am uncertain about the best checkpoint for inference. Typically, I would consider both the train loss and the validation loss to select a few optimal checkpoints for further analysis. However, the code for diffusers does not include validation loss, and I am concerned that selecting the final checkpoint might lead to overfitting issues.

Validation losses are also a slightly uncertain metric for diffusion models since we're sampling random time steps, which impacts the loss quite a bit.

In my experience, considering "FID + CLIP scores" on your validation subset is a good metric. If you want to do that quickly even CMMD could be nice. But do note that CMMD works by using 30k samples randomly drawn from the MSCOCO validation set. MSCOCO images are not of particularly high resolution. So, if your pipeline generates too high-res images, then CMMD metric might be faulty. So, it's best to have your own validation set upon which you could compute metrics like CMMD, FID, etc.