richarizardd / self-supervised-vit-path

Self-Supervised Vision Transformers Learn Visual Concepts in Histopathology (LMRL Workshop, NeurIPS 2021)

Self-Supervised Vision Transformers Learn Visual Concepts in Histopathology

NOTE: Please see our follow-up work in CVPR 2022, which further extends this repository.

Self-Supervised Vision Transformers Learn Visual Concepts in Histopathology, LMRL Workshop, NeurIPS 2021. [Workshop] [arXiv]
Richard. J. Chen, Rahul G. Krishnan

@article{chen2022self,
  title={Self-Supervised Vision Transformers Learn Visual Concepts in Histopathology},
  author={Chen, Richard J and Krishnan, Rahul G},
  journal={Learning Meaningful Representations of Life, NeurIPS 2021},
  year={2021}
}

Summary / Main Findings:

1. In head-to-head comparison of SimCLR versus DINO, DINO learns more effective pretrained representations for histopathology - likely due to 1) not needing negative samples (histopathology has lots of potential class imbalance), 2) capturing better inductive biases about the part-whole hierarchies of how cells are spatially organized in tissue.
2. ImageNet features do lag behind SSL methods (in terms of data-efficiency), but are better than you think on patch/slide-level tasks. Transfer learning with ImageNet features (from a truncated ResNet-50 after 3rd residual block) gives very decent performance using the CLAM package.
3. SSL may help mitigate domain shift from site-specific H&E stainining protocols. With vanilla data augmentations, global structure of morphological subtypes (within each class) are more well-preserved than ImageNet features via 2D UMAP scatter plots.
4. Self-supervised ViTs are able to localize cell location quite well w/o any supervision. Our results show that ViTs are able to localize visual concepts in histopathology in introspecting the attention heads.

Updates

• 06/06/2022: Please see our follow-up work in CVPR 2022, which further extends this repository.
• 03/04/2022: Reproducible and largely-working codebase that I'm satisfied with and have heavily tested.

Pre-Reqs

We use Git LFS to version-control large files in this repository (e.g. - images, embeddings, checkpoints). After installing, to pull these large files, please run:

git lfs pull

Pretrained Models

SIMCLR and DINO models were trained for 100 epochs using their vanilla training recipes in their respective papers. These models were developed on 2,055,742 patches (256 x 256 resolution at 20X magnification) extracted from diagnostic slides in the TCGA-BRCA dataset, and evaluated via K-NN on patch-level datasets in histopathology.

Note: Results should be taken-in w.r.t. to the size of dataset and duraration of training epochs. Ideally, longer training with larger batch sizes would demonstrate larger gains in SSL performance.

Arch	SSL Method	Dataset	Epochs	Dim	K-NN	Download
ResNet-50	Transfer	ImageNet	N/A	1024	0.935	N/A
ResNet-50	SimCLR	TCGA-BRCA	100	2048	0.938	Backbone
ViT-S/16	DINO	TCGA-BRCA	100	384	0.941	Backbone

Data Download + Data Preprocessing

• CRC-100K: Train and test data can be downloaded as is via this Zenodo link.
• BreastPathQ: Train and test data can be downloaded from the official Grand Challenge link.
• TCGA-BRCA: To download diagnostic WSIs (formatted as .svs files) and associated clinical metadata, please refer to the NIH Genomic Data Commons Data Portal and the cBioPortal. WSIs for each cancer type can be downloaded using the GDC Data Transfer Tool.

For CRC-100K and BreastPathQ, pre-extracted embeddings are already available and processed in ./embeddings_patch_library. See patch_extraction_utils.py on how these patch datasets were processed.

Additional Datasets + Custom Implementation: This codebase is flexible for feature extraction on a variety of different patch datasets. To extend this work, simply modify patch_extraction_utils.py with a custom Dataset Loader for your dataset. As an example, we include BCSS (results not yet updated in this work).

• BCSS (v1): You can download the BCSS dataset from the official Grand Challenge link. For this dataset, we manually developed the train and test dataset splits and labels using majority-voting. Reproducibility for the raw BCSS dataset may be not exact, but we include the pre-extracted embeddings of this dataset in ./embeddings_patch_library (denoted as version 1).

Evaluation: K-NN Patch-Level Classification on CRC-100K + BreastPathQ

Run the notebook patch_extraction.ipynb, followed by patch_evaluation.ipynb. The evaluation notebook should run "out-of-the-box" with Git LFS.

Evaluation: Slide-Level Classification on TCGA-BRCA (IDC versus ILC)

Install the CLAM Package, followed by using the 10-fold cross-validation splits made available in ./slide_evaluation/10foldcv_subtype/tcga_brca. Tensorboard train + validation logs can visualized via:

tensorboard --logdir ./slide_evaluation/results/

Visualization: Creating UMAPs

Install umap-learn (can be tricky to install if you have incompatible dependencies), followed by using the following code snippet in patch_extraction_utils.py, and is used in patch_extraction.ipynb to create Figure 4.

Visualization: Attention Maps

Attention visualizations (reproducing Figure 3) can be performed via walking through the following notebook at attention_visualization_256.ipynb.

Issues

• Please open new threads or report issues directly (for urgent blockers) to [email protected].
• Immediate response to minor issues may not be available.

Acknowledgements, License & Usage

• Part of this work was performed while at Microsoft Research. We thank the BioML group at Microsoft Research New England for their insightful feedback.
• If found our work useful in your research, please consider citing our work(s) at:

@article{chen2022self,
  title={Self-Supervised Vision Transformers Learn Visual Concepts in Histopathology},
  author={Chen, Richard J and Krishnan, Rahul G},
  journal={Learning Meaningful Representations of Life, NeurIPS 2021},
  year={2021}
}

@inproceedings{chen2022scaling,
  title={Scaling Vision Transformers to Gigapixel Images via Hierarchical Self-Supervised Learning},
  author={Chen, Richard J and Chen, Chengkuan and Li, Yicong and Chen, Tiffany Y and Trister, Andrew D and Krishnan, Rahul G and Mahmood, Faisal},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  year={2022}
}

© This code is made available under the GPLv3 License and is available for non-commercial academic purposes.

self-supervised-vit-path's Issues

Training

Hi,

Could you share the code about training too?

Thanks

Training time consumed

Hi,
Thank you for your nice work.
I wander if you could tell me training time consumed and your GPU type.

Question about loss function

Hi.

Since reading your paper, I'm quite confused about why did you use the loss function like above? or maybe the typos?
(I cannot find your train codes so only way to check is by making issues)

Loss function you used is some weird cause it measures the cross entropy of teacher distribution relative to student distribution and train student network.

Your reference paper DINO use different loss function as below, which measures the cross entropy of student distribution relative to teacher distribution and train student network.

Is there some missing points that I don't know or it's just typos?

Training data

Thanks for your great work!
In the paper, you said you used 1038 TCGA-BRCA for training. The project TCGA-BRCA provides 1062 Diagnostic Slides and 1093 Tissue Slides, and I wonder which type you used.

The two libs can't be found in the project.

from pl_bolts.models.self_supervised import resnets
from pl_bolts.utils.semi_supervised import Identity
The two libs can't be found in the project.

DINO Parameters and Code?

Hi @Richarizardd,

Super interesting work!
Would you mind providing me with the conditions/parameters you used for DINO? I am trying to replicate some parts of your paper, and would really appreciate if you could provide the code for DINO that you used. Furthermore, would you have any advice on the set of parameters I should employ when using a ResNet50 as a backbone encoder with DINO? Or have you only experimented with ViTs as the backbone?

Thanks in advance :)

global pooling in ViT backbone

Thanks for your great work.
In the figure, there is a global pooling layer after transformer encoder in ViT backbone. In the original implementation of dino, only cls token is used for successive processing. I wonder if you use global pooling of all tokens instead of cls token in this step.