print ema steps and resume v2 pretraining

This commit is contained in:
Patrick Esser 2022-08-02 20:32:22 +00:00
parent 85868a5d34
commit c229c115c1
2 changed files with 4 additions and 2 deletions

View file

@ -10,6 +10,8 @@ def printit(p):
sd = torch.load(p, map_location="cpu") sd = torch.load(p, map_location="cpu")
if "global_step" in sd: if "global_step" in sd:
print(f"This is global step {sd['global_step']}.") print(f"This is global step {sd['global_step']}.")
if "model_ema.num_updates" in sd["state_dict"]:
print(f"And we got {sd['state_dict']['model_ema.num_updates']} EMA updates.")
if __name__ == "__main__": if __name__ == "__main__":

View file

@ -24,8 +24,8 @@ cd /fsx/stable-diffusion/stable-diffusion
CONFIG=configs/stable-diffusion/v2_pretraining.yaml CONFIG=configs/stable-diffusion/v2_pretraining.yaml
# resume and set new seed to reshuffle data # resume and set new seed to reshuffle data
EXTRA="--seed 542 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/checkpoints/v2-256/216k-256.ckpt" #EXTRA="--seed 542 model.params.ckpt_path=/fsx/stable-diffusion/stable-diffusion/checkpoints/v2-256/216k-256.ckpt"
#EXTRA="--seed 543 --resume_from_checkpoint ..." EXTRA="--seed 543 --resume_from_checkpoint /fsx/stable-diffusion/stable-diffusion/logs/2022-07-31T23-35-31_v2_pretraining/checkpoints/last.ckpt"
# reduce lr a bit # reduce lr a bit
#EXTRA="${EXTRA} model.params.scheduler_config.params.f_max=[0.75]" #EXTRA="${EXTRA} model.params.scheduler_config.params.f_max=[0.75]"