How to change the diffusion steps

[ValeriaLiu]

Thank you for your excellent work!

I run this code on my own dataset about bladder tumor but I get very poor results. I guess the reason may be that my dataset is too small, resulting in overfitting. The train_loss is very small with only about 0.003 but the Dice Score is very low. So I want to use a smaller diffusion_steps.

However , when I run the segmentation_train.py with the following instruction:
python scripts/segmentation_train.py --diffusion_steps 800
the program error occured as following:

I feel very strange because when the diffusion step is reduced, the program's memory requirements for memory of the GPU should be reduced. And I want to know that have you met the similar problem or can you tell me how to run this code with a smaller diffusion steps properly?

I would appreciate if you can reply me at your convenience. Thanks a lot.

[JuliaWolleb]

Hi

Thanks for your interest. Can you maybe plot one image of your predicted segmentation map? This might help me to see where the problem is with your low dice score.
Actually, using less diffusion steps won't solve your problem. T=1000 seems to work fine with most applications. There is work that proposes T=4000 for better results, but this will take forever. Moreover, this will not reduce GPU memory consumption. So I suggest to stick with T=1000 and investigate more in what went wrong.

[ValeriaLiu]

Hi

Thanks for your interest. Can you maybe plot one image of your predicted segmentation map? This might help me to see where the problem is with your low dice score. Actually, using less diffusion steps won't solve your problem. T=1000 seems to work fine with most applications. There is work that proposes T=4000 for better results, but this will take forever. Moreover, this will not reduce GPU memory consumption. So I suggest to stick with T=1000 and investigate more in what went wrong.

Hi

Thanks for your interest. Can you maybe plot one image of your predicted segmentation map? This might help me to see where the problem is with your low dice score. Actually, using less diffusion steps won't solve your problem. T=1000 seems to work fine with most applications. There is work that proposes T=4000 for better results, but this will take forever. Moreover, this will not reduce GPU memory consumption. So I suggest to stick with T=1000 and investigate more in what went wrong.

Thank you for your reply.
The following are the five predicted segmentation masks of one bladder tumor image.

The Dice Score above is about 0.08.

I trained the model for about 51000 steps. The loss droped below 0.01 after about 5500 steps and fluctuated between 0.003-0.005 after 30000steps.

I am truly grateful for your kind assistance.

[JuliaWolleb]

Hi

This seems like it is generating random segmentation masks, and not segmentation masks that belong to your input image. Do you properly stack the input image and the noisy segmentation mask channel-wise in every step during training and sampling?

[alex-stenger]

Hi,

Actually there is an error in the code making impossible to change the number of diffusion step.

In segmentation_train.py at line 33, the schedule sampler takes an extra argument maxt=1000, that goes into the UniformSampler class in resample.py such that even if you change -diffusion_steps flag, the Sampler will everytime sample with 1000 diffusion step.

I recommand to check the original code of OpenAI

and replace the class UniformSampler by the following one :

class UniformSampler(ScheduleSampler):
    def __init__(self, diffusion):
        self.diffusion = diffusion
        self._weights = np.ones([diffusion.num_timesteps])
    def weights(self):
        return self._weights

[JuliaWolleb]

Hi
Oh sorry, yes if you change the number of diffusion steps, you also need to adapt the maxt accordingly. But this was only I parameter I played around with, so dropping maxt does not make any difference. Thanks for fixing it!