ViT not learning for binary classification problem

[bax24]

Hello everyone !

I didn't know whether to ask this in the discussion forum or to start a "New issue" so I hope this is the right place..

I am currently trying to implement a classification network with MRI scans (from ADNI) to detect alzheimer. I already implemented 3D CNN network which works quite well and I wanted to test the backbone of UNETR, so the ViT, to see what kind of results I could get.

However, the network doesn't seem to learn as the loss remains the same throughout epochs. I tried with a classification of 1 class and 2 classes and the results are the same. I suppose I am not using it correctly and I was wondering if you could help me ?

I attached the codes for my main, my model, how I train and test it. I hope this is ok I didn't know whether a github would be easier for you or not (since my code is still a bit messy).

I hope to hear from you and I thank you in advance,

Kind regards,

Lucas

[finalelement]

Hi @bax24 thanks for starting the discussion. Could you share a github link with the CNN code as well?

The ViT usage at the first glance does not seem to be incorrect, there does seem to be some confusion between out_channel and num_classes as that is not being tied together, please check and verify that.

Another thing that is interesting here is that how are you testing with a fixed input patch size. Is your CNN network also monai based?

[bax24]

Hello,

What do you mean when you say 'tied together' i'm sorry I don't get it? The sizes of my MRI is (96, 96, 48) and it only has one channel (I believe).

And for the patch size I just used (16, 16, 16) for training and testing, should I use something else ?

The CNN is not MONAI based, its implementation uses simple 3D CNN layers and fully connected ones.

Here is my github with the relevant codes : https://github.com/bax24/Alzheimer-Classification

Thanks again for your help!

PS : I am still a student and this is part of my master thesis so if you see something outrageous in the way that I am training or testing my models don't hesitate to tell me

[finalelement]

What I mean to say by 'not tied together' is here you define out_channel=1 L204 but the out_channel is not being assigned to num_classes [here]
(https://github.com/bax24/Alzheimer-Classification/blob/874b60dd0ac50ac4f6e4e91b74a4c17ca8e81ff1/model.py#L138)

I realize that you are downsampling all your images to be of the same resolution, with 96x96x48

For the patch size of ViT 16x16x16 is good and should work fine.

Can you share what you see with the CNN, the plots?

[bax24]

Here is the plot you asked for i am sorry i forgot to send it

[bax24]

The CNN models now gives very good results but I still cannot make the ViT model learn (the loss is constant throughout the learning

No matter the epoch I get the following

I have updated my github (https://github.com/bax24/Alzheimer-Classification) could you please take a look and tell me if i am doing something wrong ? Maybe with the backpropagation ?

Thanks a lot in advance,

Lucas

[finalelement]

I took a look at your code, here are suggestions which could help:

• https://github.com/bax24/Alzheimer-Classification/blob/ecd3aafc0e67c3e5a07f0e678224ae6803d009fd/model.py#L159 why is the activation not being used for the ViT part?
• Could try this with a lowered learning rate of probably 1e-4.
• I noticed the plots of training and testing/validation for the CNN model, looks quite nice, can you share these plots when you run ViT instead of CNN?

[bax24]

Also, should I rewrite (like in NetworkB in my github) the init and forward method or should I just use monai.Vit as the model itself ?

[bax24]

Here is the plot when running the ViT with a lower learning rate of 1e-4.

Everything remain constant..

I do not use the activation because I already use a Sigmoid and it didn't work either with the log_softmax do you think I should try again ?

[chris-clem]

Hi,

I noticed that the Vision Transformer (ViT) with classification=True has nn.Tanh() in the end of the classification_head (https://github.com/Project-MONAI/MONAI/blob/dev/monai/networks/nets/vit.py#L100). I did the following and it worked:

model = monai.networks.nets.ViT(...)
model.classification_head = nn.Sequential(nn.Linear(768, num_classes))

Best
Christoph

[finalelement]

@bax24 the above is worth a shot, as you mentioned that you have an activation of sigmoid addiitionally.

[bax24]

When running the model proposed by @chris-clem I get the following error and I cannot figure it out

Which arguments are they talking about ?

[bax24]

I get the same mistakes when I run the example proposed on the doc.monai.io website :

" # for 3-channel with image size of (128,128,128), 24 layers and classification backbone

net = ViT(in_channels=3, img_size=(128,128,128), pos_embed='conv', classification=True) "

After trying multiple solutions, the best I can get is the network varying for 3 or 4 epochs and then stays constant along all the rest

This plot is obtained with the updated github (https://github.com/bax24/Alzheimer-Classification)