ViT for 3D brain images

[Meddebma]

Dear all,

I tried to make a classification task of 3D brain images where I got a dice score >0.8 using DenseNet121, and I wanted to test whether I could improve it using ViT. When I tried the ViT implementation I've got this Error:
cross_entropy_loss(): argument 'input' (position 1) must be Tensor, not tuple

Could you please take a look on this notebook:
https://github.com/Meddebma/pyradiomics/blob/master/Classification_HBI_ViT.ipynb

Thank you very much

[Nic-Ma]

Hi @Meddebma ,

Thanks for your experiments and feedback here.
I think the root cause is that the ViT returns a tuple of 2 data instead of only 1 array:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/networks/nets/vit.py#L112

Hi @finalelement @yiheng-wang-nv @wyli , I plan to slightly enhance the forward() logic of ViT to make it optional to return the hidden_states_out, similar to our DynUNet implementation:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/networks/nets/dynunet.py#L243
I remember there may be some TorchScript and distributed data-parallel issue in DynUNet with the uncertain return type.
Do you have any concerns?

Thanks.

[finalelement]

@Meddebma I went through the notebook. This could also work out if you take outputs[0] for the loss instead of just outputs. since the ViT provides classification and the hidden states both. loss = loss_function(outputs, labels)

In brief, I agree with @Nic-Ma we could make the hidden states optional for easier use.

[Meddebma]

Hi @finalelement,
thanks for your help, I implemented your solution, but I've got the same problem in the validation after the second epoch, also the loss looked weired since it did not change. Here is the notebook:

https://github.com/Meddebma/pyradiomics/blob/master/Classification_HBI_ViT.ipynb

thanks

[Nic-Ma]

@finalelement , let me try to fix the TorchScript and distributed data-parallel issues and add support to optionally return the hidden_states_out ASAP.

Thanks.

[Nic-Ma]

Submitted draft PR for this ticket: Project-MONAI/MONAI#3428.

Thanks.

[finalelement]

@Meddebma you ran into that issue, because the model's prediction term also needs to be something like model(val_images)[0]

y_pred = torch.cat([y_pred, **model(val_images)**], dim=0) y = torch.cat([y, val_labels], dim=0)

I would suggest to try fixing the validation, it is not the same way as how the prediction is being extracted from the training. This should hopefully resolve the second epoch validation issue. :)

[Meddebma]

Thank you very much @finalelement, this time it worked.

Unfortunately, something is not working in the classification, I get a accuracy of 0.5 that does not improve over the epochs
Is it a problem specific for my 3D Dataset?

https://github.com/Meddebma/pyradiomics/blob/master/Classification_HBI_ViT.ipynb

[finalelement]

@Meddebma Have you tried this exact same data loading and post-processing with a smaller network architecture or any other network, to make sure everything else (all the other components) are working before concluding that there is something incorrect with ViT?

[Meddebma]

Hi @finalelement
I have tried exactly the same pipeline with Densenet212 and got an accuracy of 0.83 and AUC of 0.97,
I only changed the model.
Here is the notebook:
https://github.com/Meddebma/pyradiomics/blob/master/Classification_HBI.ipynb

[LouiseBloch]

I have the same problem with my 3D classification task and the VisionTransformer as well and no solution by now. However, if I train the pipeline using the DenseNet121 or the EfficientNetB0, it worked.

Unfortunately, something is not working in the classification, I get a accuracy of 0.5 that does not improve over the epochs Is it a problem specific for my 3D Dataset?

[Nic-Ma]

Hi @finalelement ,

Could you please help share more details to help figure out the issues?

Thanks in advance.

[finalelement]

@LouiseBloch @Meddebma are you guys specifically facing this issue with 3D classification task or is this reproducible for 2D as well?

I just verified the test cases, they look okay to me. We can perhaps resort to CIFAR testing before proceeding with this task for this net to narrow down on the issue?

[ahatamiz]

@Nic-Ma This seems to be a very straightforward issue to deal with. The output of ViT is a tuple. The first need to be used for classification. This was done by design since we make use of intermediate feature maps in the UNETR:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/networks/nets/unetr.py#L191

Furthermore, a proper classification loss need to be used.

Maybe I should do a tutorial for ViT classification to further demonstrate the concept.

[Nic-Ma]

Hi @ahatamiz ,

Thanks for your update here.
BTW, I am working on a draft PR to enhance the ViT return data, we need to make sure TorchScript and Distributed Data Parallel can work well with it.

Thanks.