filter_num in TransUNet

[parniash]

Hi,

In TransUNet, filter_num shows the number of filters for down and upsampling levels, right? However, If I use 3 filters filter_num=[64, 128, 128] instead of the default filter_num=[64, 128, 128, 512], the number of parameters of the network increases, and I get an OOM error. Is this a bug or am I missing something?

[yingkaisha]

If you do this on GPUs, then a possible reason is that your configuration is too big. [64, 128, 128] --> [64, 128, 128, 256] adds a lot of weights.

[parniash]

The problem is that I go from [64, 128, 128, 256] --> [64, 128, 128] and the network gets larger! This doesn't make sense since I reduce the number of layers but the network gets bigger.

[parniash]

Have you been able to reproduce this issue?

[yingkaisha]

@parniash Would you mind sharing your code? I don't think the network would get bigger.

[parniash]

If you compile these two models (filter_num is the difference):

model = models.transunet_2d((input_height, input_width, 3), filter_num=[32, 64, 128, 256], n_labels=n_classes, embed_dim=600, num_mlp=1000, num_heads=4, num_transformer=4, activation='ReLU', mlp_activation='GELU', output_activation='Sigmoid', batch_norm=True, pool=True, unpool='bilinear', backbone='ResNet50V2', weights='imagenet', freeze_backbone=True, freeze_batch_norm=True, name='transunet')

vs

model = models.transunet_2d((input_height, input_width, 3), filter_num=[64, 128, 256], n_labels=n_classes, embed_dim=600, num_mlp=1000, num_heads=4, num_transformer=4, activation='ReLU', mlp_activation='GELU', output_activation='Sigmoid', batch_norm=True, pool=True, unpool='bilinear', backbone='ResNet50V2', weights='imagenet', freeze_backbone=True, freeze_batch_norm=True, name='transunet')

The last one has more parameters. I don't understand why.

[yingkaisha]

If I plug in: (128, 128, 3), n_labels=2

Model 1: Total params: 30,508,410
Model 2: Total params: 29,812,794

So your second configuration is smaller, there is no problem.

I feel that you are commenting on the number of trainable parameters---the second one has more trainable params, because its output head is connected to 64 channels.

Try dig into your configurations with model.summary(). The total size of a model can be reduced in many ways.