Code for CVPR 2024 paper, "VoCo: A Simple-yet-Effective Volume Contrastive Learning Framework for 3D Medical Image Analysis"
Authors: Linshan Wu, Jiaxin Zhuang, and Hao Chen
This work presents VoCo, a simple-yet-effective contrastive learning framework for pre-training a large scale 3D medical images. Our 10k CT images pre-training model are available. Our 100k CT images pre-training model comming soon!
Self-Supervised Learning (SSL) has demonstrated promising results in 3D medical image analysis. However, the lack of high-level semantics in pre-training still heavily hinders the performance of downstream tasks. We observe that 3D medical images contain relatively consistent contextual position information, i.e., consistent geometric relations between different organs, which leads to a potential way for us to learn consistent semantic representations in pre-training. In this paper, we propose a simple-yet-effective Volume Contrast (VoCo) framework to leverage the contextual position priors for pre-training. Specifically, we first generate a group of base crops from different regions while enforcing feature discrepancy among them, where we employ them as class assignments of different regions. Then, we randomly crop sub-volumes and predict them belonging to which class (located at which region) by contrasting their similarity to different base crops, which can be seen as predicting contextual positions of different sub-volumes. Through this pretext task, VoCo implicitly encodes the contextual position priors into model representations without the guidance of annotations, enabling us to effectively improve the performance of downstream tasks that require high-level semantics. Extensive experimental results on six downstream tasks demonstrate the superior effectiveness of VoCo.
Our checkpoints 10K pre-trained checkpoint is available at VoCo_10k.pt. More results are comming.
| Method | Dataset | Pre-trained model | Training log | BTCV |
|---|---|---|---|---|
| VoCo | 10k CT | VoCo_10k.pt | Part of Pre-training log | 84.51 |
First, you need to download the pre-training dataset. The 10k dataset are all open-source and you can download yourself. Or you can download it in our hugging face repo.
Note: 10k dataset is collected by Dr. Jiaxin Zhuang
├── data
├── BTCV
├── TCIAcovid19
├── Luna16-jx
├── stoic21
├── Totalsegmentator_dataset
├── Flare23
├── LiDC
└── HNSCC_convert_v1
(1) Note that in this repo, we present the version of our 10k pre-training, thus some details may be different to our paper.
(2) To accerlate the training, we use "Persistentdataset" to pre-cache dataset, which requires extra storage. It is important in our codes. If you don't have enough storage, you can change it back in "utils/data_utils.py".
To pre-train:
sh train.shOur finetune codes will soon be available, or you can directly use the codes in MONAI.
We thank MONAI for part of their codes.
If you find this repo useful for your research, please consider citing the paper as follows:
@inproceedings{VoCo,
title={VoCo: A Simple-yet-Effective Volume Contrastive Learning Framework for 3D Medical Image Analysis},
author={Wu, Linshan and Zhuang, Jiaxin and Chen, Hao},
booktitle={IEEE Conf. Comput. Vis. Pattern Recog.},
year={2024}
}