Nature Medicine 
Journal Link | Open Access Read Link | Download Models | Download Pre-extracted Embeddings | Cite
- 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 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 |
Last Updated 3/20/2025
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.
| 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 |
| 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.
First clone the repo and cd into the directory:
git clone https://github.com/mahmoodlab/UNI.git
cd UNIThen create a conda env and install the dependencies:
conda create -n UNI python=3.10 -y
conda activate UNI
pip install -e .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.
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)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.
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_voteRefer to the notebooks below for detailed examples.
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.
ⓒ 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.
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.
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}
}
