UNI

🚀 UNIモデルの紹介

UNIは、組織病理学用の最大規模の事前学習済みビジュアルエンコーダです。内部の腫瘍、感染、炎症、正常組織に基づいて開発され、公開されています。このモデルは34の臨床タスクで優れた性能を発揮し、特にまれながんや代表が不足しているがんタイプで顕著な性能向上が見られます。

🚀 クイックスタート

モデルの更新

UNIの後継モデルであるUNI2がリリースされました。これは35万枚以上の異なる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()  # ユーザーアクセストークンでログインします。https://huggingface.co/settings/tokensで確認できます。

# UNIの重みを読み込むにはpretrained=Trueが必要です（初回は重みをダウンロードします）
# LayerScaleパラメータを正常に読み込むにはinit_valuesを渡す必要があります（例: - 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()  # ユーザーアクセストークンでログインします。https://huggingface.co/settings/tokensで確認できます。

local_dir = "../assets/ckpts/vit_large_patch16_224.dinov2.uni_mass100k/"
os.makedirs(local_dir, exist_ok=True)  # ディレクトリが存在しない場合は作成します
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) # 画像のリサイズと正規化（ImageNetのパラメータ）後の形状が[1, 3, 224, 224]の画像（torch.Tensor）
with torch.inference_mode():
    feature_emb = model(image) # 抽出された特徴（torch.Tensor）、形状は[1,1024]

直接使用（事前抽出され凍結された特徴を使用）

このモデルは、微調整せずに以下のタスクに使用して、競争力のある結果を得ることができます：

• ROI分類：クラスラベルにロジスティック回帰分類器を適用します。
• ROI分類：クラスラベルにk近傍（k - NN）分類器を適用します。
• ROI分類：クラスラベルに最近傍重心分類器（SimpleShot）を適用します。
• ROI検索：最近傍分類器を使用します。
• スライス分類：WSIから抽出されたクラスラベルのパッケージに多インスタンス学習（MIL）分類器を適用します。

下流の使用（微調整）

モデルを微調整することもでき、分割タスクでは微調整を行うことをおすすめします。ViTを密集予測タスクに適応させるための専用フレームワーク、例えばViTDetやViT - Adapter（Mask2Formerに依存）を使用して微調整することをおすすめします。

🔧 技術詳細

学習の詳細

• 学習データ：Mass - 100K、事前学習データセット（MGH、BWH、GTExから取得）。20倍の解像度で、100,402枚のH&E全スライス画像からサンプリングされた75,832,905枚の[256×256]と24,297,995枚の[512×512]の組織学画像で構成されています（合計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にも功績を認めるべきです。