How to continue the training process from a checkpoint

[LeCongThuong]

Hi @mkshing,
Your work is amazing, and I have tried to experiment with SVDiff on Google Colab with your repo.
But I actually got a little problem. Because the time for a Google Colab session is very limited, so I can not complete the training process in one Google Colab session. When I tried to resume the training process with "--resume_from_checkpoint", I coped with the following error:
FileNotFoundError: [Errno 2] No such file or directory: /path/to/checkpoints/pytorch_model.bin
If I understand correctly, your repo only saves spectral_shifts_te information, but accelerator.load_state() needs more (u-net/text encoder weight). Any solution for that?
Btw, thank you again for your repo. It is great when you publicize the repo.

[mkshing]

@LeCongThuong yeah unfortunately —resume_from_checkpoint is not supposed now and waiting until the training finishes is only you can do atm.
I’ll fix it asap.

[dome272]

Maybe this helps:

    if args.resume_from_checkpoint:
        path = os.path.basename(args.resume_from_checkpoint)

        if path is None:
            accelerator.print(
                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
            )
            args.resume_from_checkpoint = None
        else:
            # accelerator.print(f"Resuming from checkpoint {path}")
            # accelerator.load_state(os.path.join(args.output_dir, path))
            global_step = int(path.split("-")[1])

            resume_global_step = global_step * args.gradient_accumulation_steps
            first_epoch = global_step // num_update_steps_per_epoch
            resume_step = resume_global_step % (num_update_steps_per_epoch * args.gradient_accumulation_steps)

    if args.train_text_encoder:
        if args.resume_from_checkpoint:
            text_encoder, trained_token_embeds = load_text_encoder_for_svdiff(args.pretrained_model_name_or_path, 
                                                        subfolder="text_encoder", revision=args.revision, 
                                                        spectral_shifts_ckpt=args.resume_from_checkpoint)
        else:
            text_encoder, trained_token_embeds = load_text_encoder_for_svdiff(args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision)
    else:
        text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder", revision=args.revision)

    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
    if args.resume_from_checkpoint:
        unet = load_unet_for_svdiff(args.pretrained_model_name_or_path,
                                    subfolder="unet", revision=args.revision, low_cpu_mem_usage=True,
                                    spectral_shifts_ckpt=args.resume_from_checkpoint)
    else:
        unet = load_unet_for_svdiff(args.pretrained_model_name_or_path, subfolder="unet", revision=args.revision, low_cpu_mem_usage=True)

[mkshing]

Thanks @dome272 !
Could you possibly make a PR for it? Then, I will merge it.

[LeCongThuong]

Hi, @dome272
Thank you for proposing the solution. But I have 2 questions

1. Firstly, what is the meaning of the "trained_token_embeds" variable? In the original code, the load_text_encoder_for_svdiff() function only returns CLIPTextModel, why do you need that variable?
2. Secondly, after removing the "trained_token_embeds" variable, I coped with another error:
   RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and CPU! at the following line:
   encoder_hidden_states = text_encoder(batch["input_ids"])[0]
   I have not known much about the Accelerate library, but I think that the cause of the error is switching the loading models with Accelerate.load_state() into the custom load. So, I try to add ".to(accelerator.device, dtype=weight_dtype)" into Unet, TextEncoder, and AutoencoderKL but the OOM error happens. Might you have the same problem?

Thank you!

[dome272]

Sorry this took so long. I actually wanted to make a pull-request and add all my changes. But I'm quite busy atm and my codebase became very different from this one here. I've just uploaded all my stuff from my hacked repo here https://github.com/dome272/SVDiff-pytorch. Feel free to take out the necessary stuff you need for loading a checkpoint. I'm also training a separate token embedding for the object that is finetuned. I think I'm getting better results than without. But I didn't do much tests, maybe it was just luck and it's not so important.