Creating Various Checkpoint During Training

[cryptoguard]

Hey!

Is there a way to add multiple checkpoints during training? Say I want to train a model on 10,000 steps, but want to keep track of the changes every 1,000 steps without exclusively relying on the test images in the logs (i.e., by running X/Y plot charts for each model), how would I go about doing that? Any specific code I should add to the Runpod training flow?

Thanks for this awesome training protocol!

[caniyabanci76]

Yes there is a way. I've been doing runs where I have checkpoints saved at every 250 or 500 step intervals.

go to the directory Dreambooth-Stable-Diffusion/configs/stable-diffusion/

Open to edit the file v1-finetune_unfrozen.yaml

at around line 106 you will see a block :

lightning:
  modelcheckpoint:
    params:
      every_n_train_steps: 500
  callbacks:
    image_logger:
      target: main.ImageLogger
      params:
        batch_frequency: 500
        max_images: 8
        increase_log_steps: False

  trainer:
    benchmark: True
    max_steps: 3000

Change it so that it looks something like this :

lightning:
  modelcheckpoint:
    params:
      every_n_train_steps: 500      
      auto_insert_metric_name: false
      filename: "epoch{epoch:03d}-step{step:04d}-loss{val/loss_simple_ema:.4f}"
      save_top_k: 5 # saves the n number of checkpoints based on least loss
      save_last: true # can set to true if you still want the last one      
  callbacks:
    metrics_over_trainsteps_checkpoint:
      target: pytorch_lightning.callbacks.ModelCheckpoint
      params:
        every_n_train_steps: 500
        save_weights_only: True
    image_logger:
      target: main.ImageLogger
      params:
        batch_frequency: 100
        max_images: 4 
        increase_log_steps: False

  trainer:
    benchmark: True
    max_steps: 3000
    val_check_interval: 500 # loss validation check every n steps 
    log_every_n_steps: 500 # log every n steps

There are two areas where checkpoints are saved:

1. modelcheckpoint: this block refers to checkpoints saved in the directory: Dreambooth-Stable-Diffusion/logs/<session_name>/checkpoints/. It will save every 500, but will only only those based on best loss values. The number of which you can set (save_top_k).

2. This is the one you are asking for. The callbackmetrics_over_trainsteps_checkpoint . this block refers to the block in the file Dreambooth-Stable-Diffusion/main.py (line 783 at time of writing). The params here will save a checkpoint at the interval specified. and they will be saved in the directory: Dreambooth-Stable-Diffusion/logs/<session_name>/checkpoints/training_steps

Note : you may find that the folder Dreambooth-Stable-Diffusion/logs/<session_name>/checkpoints/ is locked and inaccessible via the file manager. If that is the case then you will have to use the terminal to get into it and copy or move the checkpoints to a folder where you can get to them to download them

That's it. Make sure you allocate enough drive space for your GPU instance. each checkpoint takes up about 2GB.

[cryptoguard]

This worked perfectly, thanks a ton my friend, this will save me a lot of GPU rental time!

Also, can this be made a default option within the Colab/Runpod/Vast flow? Would help a ton of people who do not know of the manual way to do it, I am sure @JoePenna

[geocine]

I followed this and works however I notice , there is ony 1 image save per 100 steps with this config

    image_logger:
      target: main.ImageLogger
      params:
        batch_frequency: 100
        max_images: 4 
        increase_log_steps: False

With this config I am expecting 4 images would be saved, but I only see 1. Do you know how I can fix this?

[caniyabanci76]

I've not paid too much attention to this to be honest. This was taken from a discord discussion to help speed up the training a bit if you don't need too many samples to see how things are progressing. you can read it here.

https://discord.com/channels/1023277529424986162/1026874778331643995/1034886537227276318

[dragoriano]

Yes there is a way. I've been doing runs where I have checkpoints saved at every 250 or 500 step intervals.

go to the directory Dreambooth-Stable-Diffusion/configs/stable-diffusion/

Open to edit the file v1-finetune_unfrozen.yaml

at around line 106 you will see a block :

lightning:
  modelcheckpoint:
    params:
      every_n_train_steps: 500
  callbacks:
    image_logger:
      target: main.ImageLogger
      params:
        batch_frequency: 500
        max_images: 8
        increase_log_steps: False

  trainer:
    benchmark: True
    max_steps: 3000

Change it so that it looks something like this :

lightning:
  modelcheckpoint:
    params:
      every_n_train_steps: 500      
      auto_insert_metric_name: false
      filename: "epoch{epoch:03d}-step{step:04d}-loss{val/loss_simple_ema:.4f}"
      save_top_k: 5 # saves the n number of checkpoints based on least loss
      save_last: true # can set to true if you still want the last one      
  callbacks:
    metrics_over_trainsteps_checkpoint:
      target: pytorch_lightning.callbacks.ModelCheckpoint
      params:
        every_n_train_steps: 500
        save_weights_only: True
    image_logger:
      target: main.ImageLogger
      params:
        batch_frequency: 100
        max_images: 4 
        increase_log_steps: False

  trainer:
    benchmark: True
    max_steps: 3000
    val_check_interval: 500 # loss validation check every n steps 
    log_every_n_steps: 500 # log every n steps

There are two areas where checkpoints are saved:

1. `modelcheckpoint:`  this block refers to checkpoints saved in the directory: `Dreambooth-Stable-Diffusion/logs/<session_name>/checkpoints/`.  It will save every 500, but will only only those based on best loss values. The number of which you can set (save_top_k).

2. This is the one you are asking for. The callback`metrics_over_trainsteps_checkpoint `.  this block refers to the block in the file `Dreambooth-Stable-Diffusion/main.py`  (line 783 at time of writing).  The params here will save a checkpoint at the interval specified. and they will be saved in the directory: `Dreambooth-Stable-Diffusion/logs/<session_name>/checkpoints/training_steps`

Note : you may find that the folder Dreambooth-Stable-Diffusion/logs/<session_name>/checkpoints/ is locked and inaccessible via the file manager. If that is the case then you will have to use the terminal to get into it and copy or move the checkpoints to a folder where you can get to them to download them

That's it. Make sure you allocate enough drive space for your GPU instance. each checkpoint takes up about 2GB.

how to convert those into several ckpt?

[OoAnikioO]

do you think is possible to modify the "Copy and name the checkpoint file" cell for copy all its checkpoints into trained models folder ?