Skip to content

Multi-head Attention-based Deep Multiple Instance Learning

License

Notifications You must be signed in to change notification settings

tueimage/MAD-MIL

Repository files navigation

MAD-MIL

Multi-head Attention-based Deep Multiple Instance Learning

alt text The attention component used in ABMIL(Left), where only one attention module is utilized, and in MAD-MIL(Right), where multiple attention modules are incorporated.

This is the PyTorch implementation of the MAD-MIL, which is based on CLAM and WSI-finetuning.

Data Preparation

For the preprocessing of TUPAC16 and TCGA datasets, we adhere to CLAM's steps, incorporating features extracted from non-overlapping 256×256 patches at 20× magnification. We share the extracted features through the following links:

  • TUPAC16
https://zenodo.org/records/10563985/files/TUPAC.zip?download=1
  • BRCA
https://zenodo.org/records/10563985/files/BRCA.zip?download=1
  • LUNG
https://zenodo.org/records/10563985/files/LUNG.zip?download=1
  • KIDNEY
https://zenodo.org/records/10563985/files/KIDNEY.zip?download=1

Training

The training can be done for different models and datasets with proper arguments like dataset_dir, task_name, model_name, lr, and reg. This is an example of training MAD-MIL on TUPAC16.

python train.py --data_root_dir feat-directory --lr 1e-4 --reg 1e-6 --seed 2021 --k 5 --k_end 5 --split_dir task_tupac16 --model_type abmil_multihead --task task_1_tumor_vs_normal --csv_path ./dataset_csv/tupac16.csv --exp_code MAD_Five_reg_1e-6  --n 5

We sweep the weight-decay value among {1e-2, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7, 1e-8} and choose the optimal value based on the validation loss.

Examples

We present attention heatmaps in the following figure to assess the interpretability of the methods.

alt text Attention heatmaps generated by different models for a Tumor slide selected from the TUPAC dataset. Top row: from left to right, the attention heatmap produced by the ABMIL and CLAM-MB. Bottom rows: from left to right, the attention heatmap generated by the MAD-MIL/3-head-1, MAD- MIL/3-head-2, MAD-MIL/3-head-3.

Reference

Please consider citing the following paper if you find our work useful for your project.

@misc{,
      title={Multi-head Attention-based Deep Multiple Instance Learning}, 
      author={},
      year={2024},
      eprint={},
      archivePrefix={},
      primaryClass={cs.CV}
}

Releases

No releases published

Packages

No packages published

Languages