DiffusionModelUNet architecture question

[Aman0phy]

I tried running the following two architectures for the DiffusionModelUNet

Model1:
model = DiffusionModelUNet( spatial_dims=2, in_channels=1, out_channels=1, num_channels=(128, 256,512,1024), attention_levels=(False, False,True, True), num_res_blocks=3, num_head_channels=256, )

Model2:
model = DiffusionModelUNet( spatial_dims=2, in_channels=1, out_channels=1, num_channels=(64,128, 256,512,1024), attention_levels=(False,False, True, True,True), num_res_blocks=3, num_head_channels=256, )

With everything else the same, 1 epoch training time for the two models are

• 6 min to run Model 1,
• 3min 20 s to run Model 2.

Given the depth of Model 2 is more (64,128,256,512,1024) vs Model 1 (128,256,512,1024)
Can anyone please explain why Model 1 occupies more GPU memory and is slower to train ?

[KumoLiu]

Hi @marksgraham, could you please take a look at this discussion and share some comments? Thanks.

[marksgraham]

Hi @Aman0phy

Each level of the UNet halves the spatial dimensions, so if your images are H*W then at level i=1,2,... they are size H/2^i * W/2^i. But your channels are increasing as you go down the levels, so roughly your activations at each level scale with H/2^i * W/2^i * channel[i]. For your choice of channels, the increase in channels per level outcompete the decrease in spatial size and you end up using more memory, because the deeper network starts on 64 channels, the less deep network on 128.