Self supervised learning model for Neuroimaging

[GuanghuiFU]

Foundation models are trained on broad data at scale and can be adapted to a wide range of downstream tasks [1]. Self-supervised learning (SSL) techniques are one of the methods to achieve this target. A good SSL feature extractor can avoid the cost of annotating large-scale datasets [2], which is the most crucial challenge in the medical image analysis domain [3]. It is a good idea to combine some SSL models with the ClinicalDL framework.

It may contain three parts, utilize data augmentation to create data for the pretext tasks, train pretext tasks for feature learning, and downstream for evaluation:

1. Data augmentation: Data augmentation techniques are the most important part of the contrastive learning approach. We could provide some basic data augmentation methods to help users to train their SSL model quickly.
2. Pretext tasks: Pretext tasks act as an important strategy for learning representation by the pseudo labels from the original data. We could provide some pretext tasks based on the augmented data.
3. Downstream tasks: To evaluate the performance of the SSL models, we could combine some models as a baseline in classification and segmentation tasks.

Recently, I am trying to apply the method in [3] to brain data, hoping to train a good self-supervised model, and build a better feature extractor based on the characteristics of the brain image.

References:
[1] Bommasani, Rishi, et al. "On the opportunities and risks of foundation models." arXiv preprint arXiv:2108.07258 (2021).
[2] Jing, Longlong, and Yingli Tian. "Self-supervised visual feature learning with deep neural networks: A survey." IEEE transactions on pattern analysis and machine intelligence (2020).
[3] Chaitanya, Krishna, et al. "Contrastive learning of global and local features for medical image segmentation with limited annotations." arXiv preprint arXiv:2006.10511 (2020).

[14thibea]

Hi @GuanghuiFU!

Thanks a lot for opening this discussion :)
I think that in the context of the current sprint code we should focus on one particular point only. Maybe you could begin with data augmentation?

Currently in ClinicaDL data augmentation is managed by the function get_transforms in utils/caps_dataset/data.py. It takes as second input a list of data augmentation transforms data_augmentation and then it translates the values given by the user in an augmentation function: for example if the user choose Noise then the function RandomNoising will be applied online to the loaded data before giving it as input to the network.

The complete list of transforms available corresponds to the dictionary defined in the function:

augmentation_dict = {
        "Noise": RandomNoising(sigma=0.1),
        "Erasing": transforms.RandomErasing(),
        "CropPad": RandomCropPad(10),
        "Smoothing": RandomSmoothing(),
        "None": None,
    }

The main problem is that these functions are not very general, and we cannot control their hyperparameters (for example RandomNoising always applies a random noise with a standard deviation between 0 and 0.1).
Maybe you could think of implemeting a more general way of doing data augmentation? For example by reusing data transforms done by TorchIO which proposed a lot of great functions to tackle data augmentation?

What do you think of this project? Would you like to contribute in this way? Or would you prefer to start with something else?

[GuanghuiFU]

Hello @14thibea ,

I think there are some SSL pretext tasks/strategies for different types of data augmentation. Such as:

1. Rotation: predict which rotation was applied [1] [2].
2. Shift: predict whether the images have been shifted.
3. Shuffle: predict the correct order of MRI or CT scan that were shuffled [5].
4. Scale: predict whether the images have been scaled.
5. Flip: predict whether the images have been flipped.
6. Random erasing/ Random masking: Reconstruct / inpainting the augmented image to the original image [3], or to predict whether the images have been masked [5].

• First, I think we could provide a data augmentation zoo which can generate data that consist with characteristics of neuroimaging. It will be more flexible that user can develop their own strategies whatever they want.

python clinicaldl --data_augmentation='rotation'

• Second, we could do data augmentation with pretext tasks. The pretext tasks can be trained by some basic neural networks. It will be easier to implement, and users can quickly train their own self-supervised learning models.

  • It could be a multi-task model that trains multiple pretext tasks together such as [5]:

    python clinicaldl SSL --pretext_task_1='rotation' --pretext_task_2='crop' --pretext_task_3='restoration' --data_augmentation_saved=True
    

  • It could be train in a constrative learning way such as [4]:

    python clinicaldl SSL --contrastive_learning_1='rotation' --contrastive_learning_1='crop' --contrastive_learning_1='shift'  
    

The TorchIO (https://torchio.readthedocs.io/transforms/transforms.html?highlight=reproducibility#reproducibility) is a good library and it has a lot of great functions. If we decided to reuse this library, we could focus on the second part: design the pretext tasks for the corresponding data augmentation strategies.

References:

[1] Zhai, Xiaohua, et al. "S4l: Self-supervised semi-supervised learning." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2019.

[2] Gidaris, Spyros, Praveer Singh, and Nikos Komodakis. "Unsupervised representation learning by predicting image rotations." arXiv preprint arXiv:1803.07728 (2018).

[3] Chen, Liang, et al. "Self-supervised learning for medical image analysis using image context restoration." Medical image analysis 58 (2019): 101539.

[4] Tian, Yonglong, Dilip Krishnan, and Phillip Isola. "Contrastive multiview coding." Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part XI 16. Springer International Publishing, 2020.

[5] Zhu, Jiuwen, et al. "Rubik’s cube+: A self-supervised feature learning framework for 3d medical image analysis." Medical image analysis 64 (2020): 101746.

[14thibea]

Actually @sophieloiz is going to work on multi-task so maybe we can reuse her project! (see discussion #255)

Now I think that we need to define more precisely your first contribution so we can help you efficiently! I see two possibilities:

• you can add the missing data augmentation procedures in ClinicaDL you mentionned in your previous command,
• you can implement the PretextManager deriving from TaskManager (in utils/task_manager/task_manager.py) which would be the module managing the pretext tasks you mentionned in the previous comment.

Do you think it could be interesting for you? If this is the case you can open the issue to precise which project interests you and then work on your fork and open a pull request to answer it 😄

[GuanghuiFU]

I think it can be divided into two parts:

1. Data augmentation: In TorchIO, it contains many data augmentation functions, we need to choose which type of augmentation is suitable for our neuroimaging. Another important thing is to define their corresponding labels. For example, for RandomFlip, the label can be 0 and 1 to indicate whether it is flipped. For many augmentation methods, TorchIO contains many parameters to support users in generating their ideal results. But for SSL, complex controllable parameters are a challenge: it may enable us to design better pretext tasks, or it may make our tasks/labels more difficult to define.

2. Stop and resume training the self-supervised learning model: the pretext tasks combination should be selected according to different tasks to build a strong SSL model. Therefore, ClinicalDL could have the function of stop, resume, and retraining. It can support users to explore their best combination.

Currently, I am doing the first part.

[14thibea]

Hello!
I don't know if you had time to begin the project.
What I suggest is that you begin by only changing the minimal amount possible of code: get_transforms in utils/caps_dataset/data.py.
You can add in this function any relevant transforms from TorchIO + suggest new arguments to manage them.
When it is done we will look at the second part!

Please open a pull request as soon as possible to check the compatibility of your changes, otherwise we may end up with the huge PR which may not be merged to ClinicaDL in the end, which may be quite disappointing.