Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Issue in using training AnimateLCM SVD #22

Open
habibian opened this issue May 3, 2024 · 21 comments
Open

Issue in using training AnimateLCM SVD #22

habibian opened this issue May 3, 2024 · 21 comments

Comments

@habibian
Copy link

habibian commented May 3, 2024

Thanks for the great work, also for releasing the training script train_svd_lcm.py.

I am trying to reproduce the results using the provided train_svd_lcm.py, but after half of the training (20,000 / 50,000 itrs) don't see any improvement neither in loss value nor generation qualities (training on a single A100 on WebVid2M).

Could you please confirm if Ishould set the hyper-params as follows?

accelerate launch train_svd_lcm.py \
--pretrained_model_name_or_path=stabilityai/stable-video-diffusion-img2vid-xt \
--per_gpu_batch_size=1 --gradient_accumulation_steps=1 \
--max_train_steps=50000 \
--width=576 \
--height=320 \
--checkpointing_steps=1000 --checkpoints_total_limit=1 \
--learning_rate=1e-6 --lr_warmup_steps=1000 \
--seed=123 \
--adam_weight_decay=1e-3 \
--mixed_precision="fp16" \
--N=40 \
--validation_steps=500 \
--enable_xformers_memory_efficient_attention \
--gradient_checkpointing \
--output_dir="outputs" \

In the current train_svd_lcm.py, the model is being trained on 576x320 resolutions, which is much lower than the standard SVD, i.e., 1024x572. Would not this cause a problem as normal (non LCM) SVD suffer from generating lower resolution videos?

Any input is much appreciated :)

@G-U-N
Copy link
Owner

G-U-N commented May 3, 2024

Hi, thanks for the interest!

  • I set the default resolution 576x 320 because I find that even trained with lower resolution, it still works to accelerate the generation of SVD. The generation quality of SVD on resolution 576 x 320 is not ideal, and therefore even trained with lower resolutions, you'd better log the generation results in the resolution of 1024x576. Also, for better results, using 1024x576 for training is definitely a better choice when the resource is sufficient.
  • Also, I would recommend slightly decreasing the learning rate since you trained the model with only one GPU.

@G-U-N
Copy link
Owner

G-U-N commented May 3, 2024

Also would say that the default hyper-parameters applied in the training script are not carefully tailored and should just be sub-optimal. For example, using EMA should generally increase the generation stability.

@habibian
Copy link
Author

habibian commented May 3, 2024

Thanks for the swift response :)

I am now switching to 4xA100 and unfortunately still see vague blobs like the attachment. Curious to know at what iterations should I expect the generations start looking like a video? :)

Thanks!

step_5000_val_img_7_2_1steps

@G-U-N
Copy link
Owner

G-U-N commented May 3, 2024

The results uploaded seem to be abnormal. It should not flash like this with unnatural colour. Here's what I obtained trained on 576x320.

Training beginning, 0-iter, cfg = 1, inference step = 4
step_1_val_img_7_1_4steps

10k iter, cfg=1, inference step = 4
step_14500_val_img_7_1_4steps

The devices are 8 A 800, and the batch size is set to 8 without gradient accumulation.

@G-U-N
Copy link
Owner

G-U-N commented May 3, 2024

I just found the code at this line was a typo, and I fixed it. Just hope it did not mislead you.

@habibian
Copy link
Author

habibian commented May 3, 2024

Amazing! It start to look good after fixing the typo.

ThanQ :)

@G-U-N
Copy link
Owner

G-U-N commented May 3, 2024

Awesome! Very glad to hear that : D.

@habibian habibian closed this as completed May 4, 2024
@habibian habibian reopened this May 5, 2024
@habibian
Copy link
Author

habibian commented May 5, 2024

Hey Fu-Yun,

After fixing the typo, I have been training the model on 8xA100s, which should be exactly like your setting then. Unfortunately, I still can't match your generations:

Training beginning, 0-iter, cfg = 1, inference step = 4
step_1_val_img_7_1_4steps

10k iter, cfg=1, inference step = 4
step_10000_val_img_7_1_4steps

20k iter, cfg=1, inference step = 4
step_20000_val_img_7_1_4steps

Any suggestion on why this is happening?

I suspect it might be from the data. Currently I am training on WebVid2M-train (results_2M_train.csv with 2.5M videos) without any particular subsampling (based on resolution, content, etc.). Could you please elaborate a bit your training data?

Also, my dataloader does not do any particular transformation/augmentation except for normalizing pixel values to [-1, 1]. Would be great if you can share your WebVid dataloader if there is any particular detail missing.

Again, thanks a lot for your great contribution :)

@G-U-N
Copy link
Owner

G-U-N commented May 5, 2024

Hey @habibian, just uploaded an example dataset.py.

In addition to that, I would recommend freezing all the convolutional layers when training because convolution layers seem to be more vulnerable for fine-tuning.

Hope this will help for better performance.

@habibian
Copy link
Author

habibian commented May 5, 2024

Thanks for the response @G-U-N .

Regarding the freezing the convolutional layers, do you mean the ones in ResBlocks? Is is part of your implementation, or I need to implement it?

Thanks!

@G-U-N
Copy link
Owner

G-U-N commented May 5, 2024

Hi @habibian,

Yes, the ResBlocks. That was not implemented in the training script. But it should be easy to achieve that through modifying this line.

@habibian
Copy link
Author

habibian commented May 6, 2024

Hey @G-U-N ,

Thanks for the input. Following your suggestion, I kept conv layers in resblocks frozen during the training as:

    for name, para in unet.named_parameters():
        # freeze resnet convs as suggested in https://github.com/G-U-N/AnimateLCM/issues/22#issuecomment-2094802365
        if 'conv' in name and not ('conv_in' in name or 'conv_out' in name):
            para.requires_grad = False
        else:
            para.requires_grad = True
            parameters_list.append(para)

I actually observe some improvements in training with this modification as:

Convs Frozen: 20k iter, cfg=1, inference step = 4
step_20000_val_img_7_1_4steps_frozen

All Finetuned: 20k iter, cfg=1, inference step = 4
step_20000_val_img_7_1_4steps

However, I still see my trained models to have much lower quality compared to the SVD checkpoint that you guys have released SVD checkpoint: . Here are some more test examples to give you some idea about how poor the quality of my replications are. So wonder if you have trained SVD checkpoint as I am doing here, or maybe there are some differences, i.e., in code, data, etc?

Thanks a lot for your guidance and support in replicating your excellent work :)

Convs Frozen: 20k iter, cfg=1, inference step = 4
step_20000_val_img_4_1_4steps
step_20000_val_img_3_1_4steps
step_20000_val_img_2_1_4steps
step_20000_val_img_1_1_4steps

@G-U-N
Copy link
Owner

G-U-N commented May 6, 2024

Hey @habibian , I would say there's no too much difference. The only difference is that I tried to freeze more weights at the beginning of training instead of fully fine-tuing. I didn't do too much ablation on that due to my limited GPU resources.

What about trying this:

 for name, para in unet.named_parameters():
    if "transformer_block" in name and "temporal_transformer_block" not in name:
        para.requires_grad = True
        parameters_list.append(para)

Again, I would recommend logging the generated videos in resolution 1024 x 576. You will not get ideal results on low resolutions even if you train the model successfully.

LMK if you get better results.

@G-U-N
Copy link
Owner

G-U-N commented May 7, 2024

Hi @habibian, just checking in to see if you have any updates. Hope everything is going well on your end!

@habibian
Copy link
Author

habibian commented May 7, 2024

Hey @G-U-N

Thanks for the suggestion and your great support here, much appreciated!

Following your last suggestion, instead of finetuning all except resblocks I am now only finetuning spatial_transformer_blocks that is actually improving the results as follows:

Finetuning all except resblocks: 20k iter, cfg=1, inference step = 4
step_20000_val_img_2_1_4steps

Finetuning spatial_transformer_blocks: 20k iter, cfg=1, inference step = 4
step_19000_val_img_2_1_4steps

Finetuning all except resblocks: 20k iter, cfg=1, inference step = 4
step_20000_val_img_3_1_4steps

Finetuning spatial_transformer_blocks: 20k iter, cfg=1, inference step = 4
step_19000_val_img_3_1_4steps

Finetuning all except resblocks: 20k iter, cfg=1, inference step = 4
step_20000_val_img_4_1_4steps

Finetuning spatial_transformer_blocks: 20k iter, cfg=1, inference step = 4
step_19000_val_img_4_1_4steps

Finetuning all except resblocks: 20k iter, cfg=1, inference step = 4
step_20000_val_img_7_1_4steps

Finetuning spatial_transformer_blocks: 20k iter, cfg=1, inference step = 4
step_19000_val_img_7_1_4steps

And, here are the 1024 x 576 generated videos using my trained checkpoint (compared to your released checkpoint):

Finetuning spatial_transformer_blocks: 20k iter, cfg=1, inference step = 4
000000

You released AnimateLCM-SVD-xt-1.1 checkpoint: ? iter, cfg=1, inference step = 4
000000

Finetuning spatial_transformer_blocks: 20k iter, cfg=1, inference step = 4
000000

You released AnimateLCM-SVD-xt-1.1 checkpoint: ? iter, cfg=1, inference step = 4
000000

As you see, there is still a gap in generation qualities, which I am not sure how can be reduced. Is the released checkpoint trained with 50K iterations? Any particular multi-stage training or lr scheduling involved?

Thanks :)

@G-U-N
Copy link
Owner

G-U-N commented May 7, 2024

Hey @habibian. Very glad to see the improvement! And I really appreciate the detailed visual ablations.

I actually conducted the training in two-stage.

  • 30k iterations with only spatial transformer block tuned with learning rate 1e-6.
if "temporal_transformer_block" not in name and "transformer_block" in name
  • 50k iterations with only temporal transformer block tuned with learning rate 3e-7. (The temporal weights of SVD is relatively large and vulnerable.)
if "temporal_transformer_block" in name

Additionally, some more iterations on larger resolutions will help enhance the performance.

Hope this will make better performance!

@habibian
Copy link
Author

habibian commented May 7, 2024

Hey @G-U-N ,

Great, thanks for the elaboration. I will follow this multi stage training and get back to you about results.

For that, could you please describe a bit the details of the large resolution training? More specifically:

  • Training videos and resolution
  • Number of iterations and learning rate

Thanks!

@G-U-N
Copy link
Owner

G-U-N commented May 7, 2024

@habibian

The details:

Training videos: bilinear interpolated webvid-2M. If you have other video dataset with larger resolution, that will be great.
Resolution: Only spatial transformer block, a 80 GB GPU should be able to train on resolution 1024x576. Only temporal transformer block, a 80 GB GPU should be able to train on resolution 768x448.
Iterations: 10k~30k
learning rate: 1e-6

@ersanliqiao
Copy link

Hey @habibian. Very glad to see the improvement! And I really appreciate the detailed visual ablations.

I actually conducted the training in two-stage.

  • 30k iterations with only spatial transformer block tuned with learning rate 1e-6.
if "temporal_transformer_block" not in name and "transformer_block" in name
  • 50k iterations with only temporal transformer block tuned with learning rate 3e-7. (The temporal weights of SVD is relatively large and vulnerable.)
if "temporal_transformer_block" in name

Additionally, some more iterations on larger resolutions will help enhance the performance.

Hope this will make better performance!

Hi, I think in stage 2, it should use unet weight saved from stage1 to initialize the unet weights of stage2, but target unet and teacher unet should be initialized from stalibity svd xt? Am I right? But the code seemed not support this??

@G-U-N
Copy link
Owner

G-U-N commented May 12, 2024

Hey @ersanliqiao.

You should load the unet and target unet from your finetuned weight and initialize the teacher unet with stability weight.

Try this at this line

from safetensors.torch import load_file
finetuned_weight = load_file("xxx.safetensors","cpu")
unet.load_state_dict(finetuned_weight)
target_unet.load_state_dict(finetuned_weight)
del finetuned_weight

@ersanliqiao
Copy link

thank you!!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

3 participants