From aa24932c517c2e2bb2e2377b80e2ca9113e3e554 Mon Sep 17 00:00:00 2001 From: tkgix Date: Sat, 10 Dec 2022 22:10:41 +0900 Subject: [PATCH] Update main.py Simple solution to a common problem with state_dict --- main.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/main.py b/main.py index 55b6ddc..5c64b7a 100644 --- a/main.py +++ b/main.py @@ -29,7 +29,15 @@ def load_model_from_config(config, ckpt, verbose=False): pl_sd = torch.load(ckpt, map_location="cpu") if "global_step" in pl_sd: print(f"ckpt: {ckpt} has {pl_sd['global_step']} steps") - sd = pl_sd["state_dict"] + + ## sd = pl_sd["state_dict"] + if "state_dict" in pl_sd: + print("load_state_dict from state_dict") + sd = pl_sd["state_dict"] + else: + print("load_state_dict from directly") + sd = pl_sd + config.model.params.ckpt_path = ckpt model = instantiate_from_config(config.model) m, u = model.load_state_dict(sd, strict=False) @@ -774,4 +782,4 @@ if __name__ == "__main__": os.rename(logdir, dst) if trainer.global_rank == 0: print("Training complete. max_steps or max_epochs reached, or we blew up.") - print(trainer.profiler.summary()) \ No newline at end of file + print(trainer.profiler.summary())