UNI

🚀 UNI模型介紹

UNI是用於組織病理學的最大預訓練視覺編碼器，基於內部腫瘤、感染、炎症和正常組織開發並公開。它在34項臨床任務中表現出色，尤其在罕見和代表性不足的癌症類型上有顯著性能提升。

🚀 快速開始

模型更新

UNI的繼任者UNI2已發佈，它在來自350k多個不同H&E和IHC切片的2億多張圖像上進行了訓練！模型權重和說明可在以下鏈接獲取：[Huggingface Repo]

模型卡片鏈接

[Journal Link] | [Open Access Read Link] | [Github Repo] | [Cite]

✨ 主要特性

UNI簡介

UNI是用於組織病理學的最大預訓練視覺編碼器（1億張圖像，10萬張全切片圖像），它基於內部腫瘤、感染、炎症和正常組織開發並公開。我們展示了它在34項臨床任務中的先進性能，尤其在罕見和代表性不足的癌症類型上有顯著性能提升。

• 為何使用UNI：UNI在預訓練時不使用開放數據集和大型公共組織學切片集合（如TCGA、CPTAC、PAIP、CAMELYON、PANDA和TCIA中的其他數據集），這些數據集通常用於計算病理學的基準測試。我們將UNI提供給研究界，用於構建和評估病理學AI模型，避免公共基準或私人組織學切片集合的數據汙染風險。

模型信息

📦 安裝指南

軟件依賴

Python包

• torch>=2.0: https://pytorch.org
• xformers>=0.0.18: https://github.com/facebookresearch/xformers
• timm>=0.9.8: https://github.com/huggingface/pytorch-image-models

倉庫

• DINOv2（自監督學習）: https://github.com/facebookresearch/dinov2
• CLAM（切片分類）: https://github.com/mahmoodlab/CLAM
• Mask2Former（細胞和組織分割）: https://github.com/facebookresearch/Mask2Former
• ViT - Adapter（細胞和組織分割）: https://github.com/czczup/ViT - Adapter
• LGSSL（線性探針和少樣本評估）: https://github.com/mbanani/lgssl

💻 使用示例

基礎用法

import timm
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)
# init_values need to be passed in to successfully load LayerScale parameters (e.g. - block.0.ls1.gamma)
model = timm.create_model("hf-hub:MahmoodLab/uni", pretrained=True, init_values=1e-5, dynamic_img_size=True)
transform = create_transform(**resolve_data_config(model.pretrained_cfg, model=model))
model.eval()

高級用法

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/vit_large_patch16_224.dinov2.uni_mass100k/"
os.makedirs(local_dir, exist_ok=True)  # create directory if it does not exist
hf_hub_download("MahmoodLab/UNI", filename="pytorch_model.bin", local_dir=local_dir, force_download=True)
model = timm.create_model(
    "vit_large_patch16_224", img_size=224, patch_size=16, init_values=1e-5, num_classes=0, dynamic_img_size=True
)
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.ToTensor(),
        transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ]
)
model.eval()

特徵提取示例

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,1024]

直接使用（使用預提取和凍結的特徵）

該模型無需微調即可用於以下任務並獲得有競爭力的結果：

• ROI分類，在類標記上應用邏輯迴歸分類器。
• ROI分類，在類標記上應用k近鄰（k - NN）分類器。
• ROI分類，在類標記上應用最近質心分類器（SimpleShot）。
• ROI檢索，使用最近鄰分類器。
• 切片分類，在從WSI提取的類標記包上應用多實例學習（MIL）分類器。

下游使用（微調）

也可以對模型進行微調，並且建議在分割任務中進行微調以獲得有競爭力的性能。我們建議使用專門用於將ViT適應於密集預測任務的框架進行微調，例如ViTDet或ViT - Adapter（依賴於Mask2Former）。

🔧 技術細節

訓練詳情

• 訓練數據：Mass - 100K，一個預訓練數據集（源自MGH、BWH和GTEx），由75,832,905張[256×256]和24,297,995張[512×512]的組織學圖像組成，分辨率為20倍，從100,402張H&E全切片圖像中採樣（總共100,130,900張圖像）。
• 訓練機制：使用PyTorch - FSDP混合精度的fp16。
• 訓練目標：採用DINOv2自監督學習配方，包含以下損失：
  • 具有多裁剪的DINO自蒸餾損失。
  • iBOT掩碼圖像建模損失。
  • [CLS]標記上的KoLeo正則化。
• 訓練時長：125,000次迭代，批量大小為3072。
• 模型架構：ViT - Large（0.3B參數）：補丁大小16，嵌入維度1024，16個頭，MLP前饋網絡。
• 使用的硬件：4x8 Nvidia A100 80GB。
• 訓練小時數：約1024 GPU小時（總共32小時）。
• 雲服務提供商：MGB ERIS研究計算核心。

📄 許可證

本模型及相關代碼根據CC - BY - NC - ND 4.0許可證發佈，僅可用於非商業、學術研究目的，並需適當引用。禁止對UNI模型及其衍生模型（包括在UNI模型輸出上訓練的模型或由UNI模型創建的數據集）進行任何商業使用、銷售或其他貨幣化操作，如需使用需事先獲得批准。下載模型需要事先在Hugging Face上註冊並同意使用條款。下載此模型即表示您同意不分發、發佈或複製該模型。如果您組織內的其他用戶希望使用UNI模型，他們必須作為個人用戶註冊並同意遵守使用條款。用戶不得嘗試重新識別用於開發基礎模型的已去標識化數據。如果您是商業實體，請聯繫相應作者。

聯繫信息

如有任何其他問題或意見，請聯繫Faisal Mahmood (faisalmahmood@bwh.harvard.edu)、Richard J. Chen (richardchen@g.harvard.edu)、Tong Ding (tong_ding@g.harvard.edu)或Ming Y. Lu (mlu16@bwh.harvard.edu)。

致謝

本項目基於一些優秀的倉庫構建，如ViT、DINOv2、LGSSL和Timm（ViT模型實現）。感謝這些作者和開發者的貢獻。

BibTeX

如果您在研究中發現我們的工作有用，請考慮引用我們的工作：

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}
}

使用UNI的工作也應歸功於ViT和DINOv2。