Missing/Unexpected Keys when load a checkpoint
viyjy opened this issue · comments
Hi, to do the stage 2 training of 2.0-base, I am using this yaml file SD-2-base-512.yam. However, this yaml file doesn't load the checkpoint from stage 1. I add a new line under trainer to handle this issue.
trainer:
load_path: sd2.0-base-256/ep0-ba550000-rank0.pt
However, I get the following error:
Found these missing keys in the checkpoint: vae.encoder.mid_block.attentions.0.to_q.weight, vae.encoder.mid_block.attentions.0.to_q.bias, vae.encoder.mid_block.attentions.0.to_k.weight, vae.encoder.mid_block.attentions.0.to_k.bias, vae.encoder.mid_block.attentions.0.to_v.weight, vae.encoder.mid_block.attentions.0.to_v.bias, vae.encoder.mid_block.attentions.0.to_out.0.weight, vae.encoder.mid_block.attentions.0.to_out.0.bias, vae.decoder.mid_block.attentions.0.to_q.weight, vae.decoder.mid_block.attentions.0.to_q.bias, vae.decoder.mid_block.attentions.0.to_k.weight, vae.decoder.mid_block.attentions.0.to_k.bias, vae.decoder.mid_block.attentions.0.to_v.weight, vae.decoder.mid_block.attentions.0.to_v.bias, vae.decoder.mid_block.attentions.0.to_out.0.weight, vae.decoder.mid_block.attentions.0.to_out.0.bias
mosaic/0 [0]:[2023-06-19 08:43:26,877][composer.core.state][WARNING] - Found these unexpected keys in the checkpoint: vae.encoder.mid_block.attentions.0.query.weight, vae.encoder.mid_block.attentions.0.query.bias, vae.encoder.mid_block.attentions.0.key.weight, vae.encoder.mid_block.attentions.0.key.bias, vae.encoder.mid_block.attentions.0.value.weight, vae.encoder.mid_block.attentions.0.value.bias, vae.encoder.mid_block.attentions.0.proj_attn.weight, vae.encoder.mid_block.attentions.0.proj_attn.bias, vae.decoder.mid_block.attentions.0.query.weight, vae.decoder.mid_block.attentions.0.query.bias, vae.decoder.mid_block.attentions.0.key.weight, vae.decoder.mid_block.attentions.0.key.bias, vae.decoder.mid_block.attentions.0.value.weight, vae.decoder.mid_block.attentions.0.value.bias, vae.decoder.mid_block.attentions.0.proj_attn.weight, vae.decoder.mid_block.attentions.0.proj_attn.bias
For the stage 2 training, do we need to set load_weights_only=True? Thanks.
Hi @viyjy, are you using the same version of huggingface diffusers to train and to load weights? It's possible that some of the modules were renamed in that library. If you want to use that exact checkpoint without changing the diffusers version, renaming th ekeys would likely work. query -> to_q, key -> to_k and so on.
@A-Jacobson Thanks. This problem is solved. Currently the major issue is #36, can you help to check it? Many thanks!!!