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。