UNI
March 26, 2025 · View on GitHub
Towards a General-Purpose Foundation Model for Computational Pathology
Nature Medicine 
Journal Link | Open Access Read Link | Download Models | Download Pre-extracted Embeddings | Cite
Updates
- 3/20/2025: One year overview of UNI & CONCH written by our team with updated table of research applications.
- 3/6/2025: Blog Post from Meta AI on our development of UNI using DINOv2.
- 01/14/2025: Release of UNI 2 trained on over 200 million pathology H&E and IHC images sampled from 350+ thousand diverse whole slide images. UNI 2 model weights, benchmark results and 25k+ pre-extracted WSI embeddings from TCGA,CPTAC, and PANDA are released.
- 12/17/2024: Research Highlight from Nature Medicine on UNI & CONCH for clinical oncology
- 03/19/2024: UNI is published! Model weights and initial benchmark results are released.
Unfamiliar with UNI? Please refer to the original README (here) for more details or refer to the accompanying Nature Medicine study (here).
Model weights
| Model Name | Release Date | Model Architecture | Download Link |
|---|---|---|---|
| UNI2-h | 01/2025 | ViT-h/14-reg8 | HF Link |
| UNI | 03/2024 | ViT-l/16 | HF Link |
Research Applications using UNI & CONCH
Last Updated 3/20/2025
Pre-extracted Embeddings
To facilitate downstream tasks, we provide pre-extracted embeddings for the UNI 2 model (UNI2-h) for TCGA, CPTAC and PANDA, which can be downloaded here.
Benchmarking UNI 2
ROI Benchmarks
| Model name | Pretraining | Model size | HEST (Regression, Public) | CRC-100K-Raw (9 classes, Public) | TCGA Uniform Tumor (32 classes, Public) | C17-WILDS (2 classes, Public) | Kather MSI （2 classes, Public) |
|---|---|---|---|---|---|---|---|
| UNI | Vision | ViT-l/16 | 0.386 | 0.925 | 0.595 | 0.972 | 0.679 |
| UNI2-h | Vision | ViT-h/14 | 0.414 | 0.957 | 0.675 | 0.977 | 0.722 |
| Virchow 2 | Vision | ViT-h/14 | 0.398 | 0.952 | 0.620 | 0.975 | 0.713 |
| Virchow | Vision | ViT-h/14 | 0.398 | 0.919 | 0.544 | 0.977 | 0.670 |
| UNI2-g-preview | Vision | ViT-g/14 | 0.416 | 0.949 | 0.690 | 0.985 | 0.725 |
| h-optimus | Vision | ViT-g/14 | 0.415 | 0.930 | 0.647 | 0.970 | 0.707 |
| Prov-GigaPath | Vision | ViT-g/14 | 0.385 | 0.929 | 0.593 | 0.961 | 0.693 |
| CONCH | Vision-language | ViT-b/16 | 0.371 | 0.941 | 0.556 | 0.967 | 0.685 |
| MUSK | Vision-language | ViT-l/16 | 0.346 | 0.913 | 0.464 | 0.954 | 0.666 |
Slide Benchmarks
| Model name | Pretraining | Model size | EBRAINS (30 classes, Public) | PANDA (5 classes, Public) | IHC ER / PR Assess. (6 classes, Internal) |
|---|---|---|---|---|---|
| UNI | Vision | ViT-l/16 | 0.682 | 0.944 | 0.776 |
| UNI2-h | Vision | ViT-h/14 | 0.711 | 0.946 | 0.794 |
| Virchow 2 | Vision | ViT-h/14 | 0.691 | 0.931 | 0.808 |
| Virchow | Vision | ViT-h/14 | 0.681 | 0.946 | 0.756 |
| UNI2-g-preview | Vision | ViT-g/14 | 0.746 | 0.953 | 0.795 |
| h-optimus | Vision | ViT-g/14 | 0.726 | 0.953 | 0.761 |
| Prov-GigaPath | Vision | ViT-g/14 | 0.687 | 0.944 | 0.775 |
| CONCH | Vision-language | ViT-b/16 | 0.689 | 0.934 | 0.794 |
| MUSK | Vision-language | ViT-l/16 | 0.660 | 0.923 | 0.764 |
In each task, for each model, we sweep over 3 learning rates (1e-5, 5e-5, 1e-4) and report the test performance corresponding to the best performing model on the validation set.
For all assessments, all models are evaluated using the global representation (e.g. CLS token) without test time augmentation.
Installation
First clone the repo and cd into the directory:
git clone https://github.com/mahmoodlab/UNI.git
cd UNI
Then create a conda env and install the dependencies:
conda create -n UNI python=3.10 -y
conda activate UNI
pip install -e .
1. Getting access
Request access to the model weights from the Huggingface model page using links provided in the Model Weights section. You will need to login to Huggingface to download the model weights.
2. Downloading weights + Creating model
Following authentication (using huggingface_hub), the pretrained checkpoints and image transforms for UNI can be directly loaded using the timm library. This method automatically downloads the model weights to the huggingface_hub cache in your home directory, which timm will automatically find when using the commands below:
import timm
import torch
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform
from huggingface_hub import login
login() # login with your User Access Token, found at https://huggingface.co/settings/tokens
# pretrained=True needed to load UNI weights (and download weights for the first time)
# using UNI2-h as example
timm_kwargs = {
'img_size': 224,
'patch_size': 14,
'depth': 24,
'num_heads': 24,
'init_values': 1e-5,
'embed_dim': 1536,
'mlp_ratio': 2.66667*2,
'num_classes': 0,
'no_embed_class': True,
'mlp_layer': timm.layers.SwiGLUPacked,
'act_layer': torch.nn.SiLU,
'reg_tokens': 8,
'dynamic_img_size': True
}
model = timm.create_model("hf-hub:MahmoodLab/UNI2-h", pretrained=True, **timm_kwargs)
transform = create_transform(**resolve_data_config(model.pretrained_cfg, model=model))
model.eval()
You can also download the model weights to a specified checkpoint location in your local directory. The timm library is still used for defining the model architecture (e.g. custom ViT-H/14). Pretrained weights and image transforms for UNI need to be manually loaded and defined.
import os
import torch
from torchvision import transforms
import timm
from huggingface_hub import login, hf_hub_download
login() # login with your User Access Token, found at https://huggingface.co/settings/tokens
local_dir = "../assets/ckpts/uni2-h/"
os.makedirs(local_dir, exist_ok=True) # create directory if it does not exist
hf_hub_download("MahmoodLab/UNI2-h", filename="pytorch_model.bin", local_dir=local_dir, force_download=True)
timm_kwargs = {
'model_name': 'vit_giant_patch14_224',
'img_size': 224,
'patch_size': 14,
'depth': 24,
'num_heads': 24,
'init_values': 1e-5,
'embed_dim': 1536,
'mlp_ratio': 2.66667*2,
'num_classes': 0,
'no_embed_class': True,
'mlp_layer': timm.layers.SwiGLUPacked,
'act_layer': torch.nn.SiLU,
'reg_tokens': 8,
'dynamic_img_size': True
}
model = timm.create_model(**timm_kwargs)
model.load_state_dict(torch.load(os.path.join(local_dir, "pytorch_model.bin"), map_location="cpu"), strict=True)
transform = transforms.Compose(
[
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
]
)
model.eval()
The function get_encoder performs the commands above, downloading in the checkpoint in the ./assets/ckpts/ relative path of this GitHub repository.
from uni import get_encoder
model, transform = get_encoder(enc_name='uni2-h', device=device)
3. Running Inference
You can use the UNI pretrained encoder to extract features from histopathology ROIs, as follows:
from PIL import Image
image = Image.open("uni.jpg")
image = transform(image).unsqueeze(dim=0) # Image (torch.Tensor) with shape [1, 3, 224, 224] following image resizing and normalization (ImageNet parameters)
with torch.inference_mode():
feature_emb = model(image) # Extracted features (torch.Tensor) with shape [1, 1536]
These pre-extracted features can then be used ROI classification (via linear probing), slide classification (via multiple instance learning), and other machine learning settings.
Overview of specific usages
We provide high-level functions for loading the model and using it for inference. For model loading, the function get_encoder performs the commands above in Step 2, downloading in the checkpoint in the ./assets/ckpts/ relative path of this GitHub repository.
from uni import get_encoder
model, transform = get_encoder(enc_name='uni2-h', device=device)
For inference:
from uni.downstream.extract_patch_features import extract_patch_features_from_dataloader
from uni.downstream.eval_patch_features.linear_probe import eval_linear_probe
from uni.downstream.eval_patch_features.fewshot import eval_knn, eval_fewshot
from uni.downstream.eval_patch_features.protonet import ProtoNet, prototype_topk_vote
Refer to the notebooks below for detailed examples.
More detailed starter code for loading / using the model:
See ./notebooks/uni_walkthrough.ipynb to get started with loading and using the model to create embeddings, and example code for extracting ROI features and performing ROI classification / retrieval.
License and Terms of Tuse
ⓒ Mahmood Lab. The models and associated code are released under the CC-BY-NC-ND 4.0 license and may only be used for non-commercial, academic research purposes with proper attribution. Any commercial use, sale, or other monetization of the UNI models and their derivatives, which include models trained on outputs from the UNI models or datasets created from the UNI models, is prohibited and requires prior approval. Downloading the model requires prior registration on Hugging Face and agreeing to the terms of use. By downloading the models, you agree not to distribute, publish or reproduce a copy of the models. If another user within your organization wishes to use the UNI models, they must register as an individual user and agree to comply with the terms of use. Users may not attempt to re-identify the deidentified data used to develop the underlying models. If you are a commercial entity, please contact the corresponding author or Mass General Brigham Innovation Office.
Acknowledgements
The project was built on top of amazing repositories such as ViT, DINOv2, LGSSL, and Timm (ViT model implementation). We thank the authors and developers for their contribution.
Reference
If you find our work useful in your research or if you use parts of this code please consider citing our paper:
Chen, R.J., Ding, T., Lu, M.Y., Williamson, D.F.K., et al. Towards a general-purpose foundation model for computational pathology. Nat Med (2024). https://doi.org/10.1038/s41591-024-02857-3
@article{chen2024uni,
title={Towards a General-Purpose Foundation Model for Computational Pathology},
author={Chen, Richard J and Ding, Tong and Lu, Ming Y and Williamson, Drew FK and Jaume, Guillaume and Chen, Bowen and Zhang, Andrew and Shao, Daniel and Song, Andrew H and Shaban, Muhammad and others},
journal={Nature Medicine},
publisher={Nature Publishing Group},
year={2024}
}
<img src=.github/joint_logo.jpg>