mueller-franzes / medfusion

Implementation of Medfusion - A latent diffusion model for medical image synthesis.
MIT License
165 stars 32 forks source link

Can't sample using pre-trained models #8

Open djFatNerd opened 1 year ago

djFatNerd commented 1 year ago

Hi, I am trying to use sample.py to test the generation results and I used these pretrained weights from your Hugginface page:

runs/2022_12_02_151805_chest_vaegan/last_vae.ckpt

runs/chest_diffusion/last.ckpt

and I got the following error:

image

May I ask you what would be possible cause for this error? Would it be the channels of the VAE doesn't match the pretrain model?

Thank you so much!

djFatNerd commented 1 year ago

Same thing happens to the patho diffusion model:

I was using these two weights: image

and this is the error: image

mueller-franzes commented 1 year ago

Hi, Unfortunately, I can't reproduce your error message. When I look at the output of print(pipeline.latent_embedder.out_enc) I get the following:

Sequential(
  (0): BasicBlock(
    (conv): Conv2d(512, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  )
  (1): BasicBlock(
    (conv): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1))
  )
) 

Regarding the second problem, what input size are you using? Please make sure you call the function with the correct input size. pipeline.sample(n_samples, (C, H, W), ...)

Unfortunately, I did not take into account when programming that it would have been better to save this in the checkpoint file. Use the following: Chest: (8, 32, 32) Colon: (4, 64, 64) Eye: (4, 32, 32)

djFatNerd commented 1 year ago

Thanks, I was able to sample after changing size to (4,64,64).

Meanwhile, may I ask you for some suggestions? I think your project and model is so great and I wanted to try it on my own dataset, but I am relatively new on diffusion and got hard stuck.

Basically, I want to try fine-tuning your pretrained patho model with some other histopathology images (256x256) but I don't know where should I start, even after reading your whole project for some days...

djFatNerd commented 1 year ago

BTW, may I ask what computing resource is required to train your full model? Thank you so much for your help!

mueller-franzes commented 1 year ago

It's great to hear that it's working now! It's hard to say in general where you should start. Generally, you need to create a class to load your data first. I would recommend you to write a small test to load some images, similar to this (tests/dataset/test_dataset.py) . Then you have to replace the dataset in train_diffusion.py with your own and use line 134 or 135 to load the old model.

I used about 24GB GPU memory (batch size = 32) - for batch_size = 1 you still need about 6GB.

djFatNerd commented 1 year ago

Thank you so much ~ I'll try these!

Yicijiuhaobala commented 1 year ago

Thanks, I was able to sample after changing size to (4,64,64).

Meanwhile, may I ask you for some suggestions? I think your project and model is so great and I wanted to try it on my own dataset, but I am relatively new on diffusion and got hard stuck.

Basically, I want to try fine-tuning your pretrained patho model with some other histopathology images (256x256) but I don't know where should I start, even after reading your whole project for some days...

@djFatNerd hello, I changed size to (4, 64,64) , but I still get a programming error, have you ever encountered it 1