| """Exports the models of the demos (examples/models), one ONNX file each. |
| |
| Every model takes `image` (1, 3, S, S), RGB in [0, 1] (S = 224 unless SIZES says otherwise; the |
| model's own normalization is part of the graph), and returns: |
| features (1, C, H, W) the last feature map, before pooling (ViTs: the patch |
| tokens after the final norm, as a grid) |
| cell_logits (1, K, H, W) classifiers only: the classification head applied to every cell (the |
| dense readout); for models that pool their features |
| (CNNs, Swin, MaxViT) the mean over the cells equals |
| the model's logits; for ViTs, which classify from |
| the class token, it is an approximation |
| logits (1, K) only for models whose prediction is not the mean of the cells (CLIP |
| zero-shot: attention pooling or a class token): the model's own logits |
| Weights are stored as float16 and cast to float32 when the model is loaded, which |
| halves the download without changing the computation. Each model is the largest |
| of its family whose file stays under 100 MB, except where a specific model is the point |
| (CLIP RN50x4: the model of the typographic attack of Goh et al. 2021). |
| |
| Writes NAME.onnx to the current folder and checks it against PyTorch on the .jpg images in $SAMPLES. |
| Run one model per process, with a memory cap: |
| |
| SAMPLES=path/to/images systemd-run --user --scope -p MemoryMax=12G uv run --no-project \ |
| --with torch --with torchvision --with timm --with onnx --with onnxruntime --with pillow \ |
| --with open_clip_torch --with safetensors --with huggingface_hub \ |
| python export_models.py NAME |
| |
| The files are published on Hugging Face (karay/demo-models), not in the blog repo. |
| """ |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnx |
| import onnxruntime as ort |
| import timm |
| import torch |
| from onnx import helper, numpy_helper |
| from PIL import Image |
| from torch import nn |
|
|
| OUT = Path.cwd() |
| |
| SAMPLES = Path(os.environ.get('SAMPLES', Path(__file__).resolve().parents[2] / 'assets' / 'img' / 'samples')) |
|
|
|
|
| |
|
|
| def cnn_cells(model, x): |
| return model.forward_features(x).permute(0, 2, 3, 1) |
|
|
|
|
| def swin_cells(model, x): |
| return model.forward_features(x) |
|
|
|
|
| def vit_cells(model, x): |
| tokens = model.forward_features(x)[:, model.num_prefix_tokens:] |
| return tokens.reshape(tokens.shape[0], *model.patch_embed.grid_size, -1) |
|
|
|
|
| |
|
|
| def linear_head(model): |
| """ResNet, Swin, ViT: the head is one Linear layer, applied to every cell.""" |
| return model.get_classifier() |
|
|
|
|
| def norm_head(model): |
| """ConvNeXt: LayerNorm after pooling, then a Linear layer.""" |
| return lambda cells: model.head.fc(pooled_layernorm(cells, model.head.norm)) |
|
|
|
|
| def norm_tanh_head(model): |
| """MaxViT: LayerNorm after pooling, Linear + tanh (linearized at the image's own |
| pre-activation, so that the mean of the cells is still tanh of it), then a Linear layer.""" |
| def readout(cells): |
| a = model.head.pre_logits.fc(pooled_layernorm(cells, model.head.norm)) |
| a_bar = a.mean(dim=(1, 2), keepdim=True) |
| t = torch.tanh(a_bar) |
| return model.head.fc(t + (1 - t ** 2) * (a - a_bar)) |
| return readout |
|
|
|
|
| def pooled_layernorm(cells, ln): |
| """Per-cell share of ``ln(mean of the cells)``: each cell centred on its own mean and |
| divided by the spread of the pooled vector, then the affine.""" |
| pooled = cells.mean(dim=(1, 2)) |
| sigma = torch.sqrt(pooled.var(dim=-1, unbiased=False) + ln.eps) |
| return (cells - cells.mean(dim=-1, keepdim=True)) / sigma[:, None, None, None] * ln.weight + ln.bias |
|
|
|
|
| |
|
|
| def timm_model(checkpoint, **kwargs): |
| return lambda: timm.create_model(checkpoint, pretrained=True, **kwargs) |
|
|
|
|
| def dino_resnet50(): |
| """DINOv1 ResNet-50: the torchvision trunk with DINO's weights (the checkpoint has no classifier).""" |
| import torchvision |
| m = torchvision.models.resnet50(weights=None) |
| url = 'https://dl.fbaipublicfiles.com/dino/dino_resnet50_pretrain/dino_resnet50_pretrain.pth' |
| m.load_state_dict(torch.hub.load_state_dict_from_url(url, map_location='cpu'), strict=False) |
| trunk = nn.Sequential(*list(m.children())[:-2]) |
| trunk.forward_features = trunk.forward |
| trunk.pretrained_cfg = {'mean': (0.485, 0.456, 0.406), 'std': (0.229, 0.224, 0.225)} |
| return trunk |
|
|
|
|
| |
|
|
| def kaiko_vits16(): |
| """Kaiko ViT-S/16 (Kaiko non-commercial licence), normalized with mean = std = 0.5.""" |
| m = torch.hub.load('kaiko-ai/towards_large_pathology_fms', 'vits16', trust_repo=True) |
| m.pretrained_cfg = {'mean': (0.5, 0.5, 0.5), 'std': (0.5, 0.5, 0.5)} |
| return m |
|
|
|
|
| def ctranspath(): |
| """CTransPath (GPL-3.0): a Swin-T with a convolutional stem, timm port by 1aurent.""" |
| from timm.layers.helpers import to_2tuple |
|
|
| class ConvStem(nn.Module): |
| def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=768, norm_layer=None, **kwargs): |
| super().__init__() |
| self.img_size, self.patch_size = to_2tuple(img_size), to_2tuple(patch_size) |
| self.grid_size = (self.img_size[0] // self.patch_size[0], self.img_size[1] // self.patch_size[1]) |
| self.num_patches = self.grid_size[0] * self.grid_size[1] |
| stem, input_dim, output_dim = [], 3, embed_dim // 8 |
| for _ in range(2): |
| stem += [nn.Conv2d(input_dim, output_dim, 3, 2, 1, bias=False), nn.BatchNorm2d(output_dim), nn.ReLU(inplace=True)] |
| input_dim, output_dim = output_dim, output_dim * 2 |
| stem.append(nn.Conv2d(input_dim, embed_dim, kernel_size=1)) |
| self.proj = nn.Sequential(*stem) |
| self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() |
|
|
| def forward(self, x): |
| return self.norm(self.proj(x).permute(0, 2, 3, 1)) |
|
|
| m = timm.create_model('hf-hub:1aurent/swin_tiny_patch4_window7_224.CTransPath', embed_layer=ConvStem, pretrained=True) |
| cfg = timm.data.resolve_model_data_config(m) |
| m.pretrained_cfg = {'mean': cfg['mean'], 'std': cfg['std']} |
| return m |
|
|
|
|
| class TorchvisionClassifier(nn.Module): |
| """A torchvision ResNet as trunk + fc, with the timm-style methods used here.""" |
|
|
| def __init__(self, resnet, mean, std): |
| super().__init__() |
| self.trunk, self.fc = nn.Sequential(*list(resnet.children())[:-2]), resnet.fc |
| self.pretrained_cfg = {'mean': tuple(mean), 'std': tuple(std)} |
|
|
| def forward_features(self, x): |
| return self.trunk(x) |
|
|
| def get_classifier(self): |
| return self.fc |
|
|
| def forward(self, x): |
| return self.fc(self.trunk(x).mean(dim=(2, 3))) |
|
|
|
|
| def tumour_detector(): |
| """Breast tumour detector, ResNet-34 on TCGA-BRCA (Le et al. 2020; CC-BY-4.0): classes [Other, Tumor]; |
| trained on 350 px tiles at 0.25 um/px resized to 224 px.""" |
| import json |
| import torchvision |
| from huggingface_hub import hf_hub_download |
| from safetensors.torch import load_file |
| repo = 'kaczmarj/breast-tumor-resnet34.tcga-brca' |
| cfg = json.load(open(hf_hub_download(repo, 'config.json'))) |
| norm = next(t for t in cfg['transform'] if t['name'] == 'Normalize')['arguments'] |
| resnet = torchvision.models.resnet34(num_classes=2) |
| resnet.load_state_dict(load_file(hf_hub_download(repo, 'model.safetensors'))) |
| return TorchvisionClassifier(resnet, norm['mean'], norm['std']) |
|
|
|
|
| |
| |
| |
|
|
| TISSUE_PROMPTS = {'tumour': 'tumor cells', 'stroma': 'stroma', 'lymphocytes': 'lymphocytes', |
| 'necrosis': 'necrosis', 'fat': 'adipose tissue'} |
| TISSUE_TEMPLATES = ['a histopathology image of {}.', 'an H&E image of {}.'] |
|
|
|
|
| def text_head(clip, tokenizer, names, templates, cache=None): |
| """Zero-shot classifier: the mean of the normalized text embeddings of the templates, per class. Computed |
| on the GPU if there is one (80 templates x 1000 classes take hours on a CPU) and cached in `cache`.""" |
| if cache and Path(cache).exists(): |
| return torch.load(cache) |
| dev = 'cuda' if torch.cuda.is_available() else 'cpu' |
| clip = clip.to(dev) |
| with torch.no_grad(): |
| rows = [] |
| for name in names: |
| t = clip.encode_text(tokenizer([tpl.format(name) if isinstance(tpl, str) else tpl(name) for tpl in templates]).to(dev)) |
| t = t / t.norm(dim=-1, keepdim=True) |
| rows.append(t.mean(0) / t.mean(0).norm()) |
| head = torch.stack(rows).cpu() |
| clip.cpu() |
| if cache: |
| torch.save(head, cache) |
| return head |
|
|
|
|
| class ClipRN(nn.Module): |
| """OpenAI CLIP ResNet (attention pooling) with its zero-shot ImageNet classifier.""" |
|
|
| def __init__(self, name='RN50x4'): |
| super().__init__() |
| import open_clip |
| from open_clip.zero_shot_metadata import IMAGENET_CLASSNAMES, OPENAI_IMAGENET_TEMPLATES |
| clip, _, pre = open_clip.create_model_and_transforms(name, pretrained='openai') |
| clip.eval() |
| self.visual = clip.visual |
| self.register_buffer('text', text_head(clip, open_clip.get_tokenizer(name), IMAGENET_CLASSNAMES, OPENAI_IMAGENET_TEMPLATES, |
| Path(__file__).resolve().parent / f'zeroshot_{name}_imagenet.pt')) |
| self.scale = float(clip.logit_scale.exp()) |
| n = pre.transforms[-1] |
| self.pretrained_cfg = {'mean': tuple(n.mean), 'std': tuple(n.std)} |
|
|
| def forward_features(self, x): |
| v = self.visual |
| return v.layer4(v.layer3(v.layer2(v.layer1(v.stem(x))))) |
|
|
| def classify(self, emb): |
| return self.scale * (emb / emb.norm(dim=-1, keepdim=True)) @ self.text.T |
|
|
| def cell_logits(self, cells): |
| pool = self.visual.attnpool |
| B, H, W, C = cells.shape |
| x = cells + pool.positional_embedding[1:].reshape(H, W, C) |
| return self.classify(pool.c_proj(pool.v_proj(x))) |
|
|
| def forward(self, x): |
| return self.classify(self.visual.attnpool(self.forward_features(x))) |
|
|
|
|
| class ClipViT(nn.Module): |
| """An open_clip ViT (class token) with a zero-shot head over tissue classes (QuiltNet-B-32, MIT).""" |
|
|
| def __init__(self, hub='hf-hub:wisdomik/QuiltNet-B-32'): |
| super().__init__() |
| import open_clip |
| clip, _, pre = open_clip.create_model_and_transforms(hub) |
| clip.eval() |
| self.visual = clip.visual |
| self.register_buffer('text', text_head(clip, open_clip.get_tokenizer(hub), list(TISSUE_PROMPTS.values()), TISSUE_TEMPLATES)) |
| self.scale = float(clip.logit_scale.exp()) |
| n = pre.transforms[-1] |
| self.pretrained_cfg = {'mean': tuple(n.mean), 'std': tuple(n.std)} |
|
|
| def tokens(self, x): |
| v = self.visual |
| x = v.conv1(x) |
| B, C, H, W = x.shape |
| x = x.reshape(B, C, H * W).permute(0, 2, 1) |
| x = torch.cat([v.class_embedding.expand(B, 1, C), x], dim=1) + v.positional_embedding |
| x = v.transformer(v.ln_pre(x)) |
| return v.ln_post(x), H, W |
|
|
| def forward_features(self, x): |
| t, H, W = self.tokens(x) |
| return t[:, 1:].reshape(t.shape[0], H, W, -1) |
|
|
| def classify(self, emb): |
| return self.scale * (emb / emb.norm(dim=-1, keepdim=True)) @ self.text.T |
|
|
| def cell_logits(self, cells): |
| return self.classify(cells @ self.visual.proj) |
|
|
| def forward(self, x): |
| t, _, _ = self.tokens(x) |
| return self.classify(t[:, 0] @ self.visual.proj) |
|
|
|
|
| class HibouB(nn.Module): |
| """Hibou-B (HistAI, Apache-2.0; DINOv2 ViT-B/14 with 4 registers; hibou_port.py) with a tissue head trained |
| on its mean patch token (hibou_b_tissue_head.pt, from gap_mil_revision/tmp/kuleuven_talk/deck/hibou_head.py): |
| classes tumour, stroma, lymphocytes, necrosis, fat; BCSS tiles of 112 um at 0.5 um/px; no lobular case in |
| training. The head is linear on the mean of the patch tokens, so the mean of the cell logits is its logits.""" |
|
|
| def __init__(self): |
| super().__init__() |
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| from hibou_port import hibou_b |
| self.vit = hibou_b() |
| saved = torch.load(Path(__file__).resolve().parent / 'hibou_b_tissue_head.pt') |
| self.head = nn.Linear(self.vit.embed_dim, len(saved['classes'])) |
| self.head.load_state_dict(saved['state_dict']) |
| self.pretrained_cfg = {'mean': tuple(saved['mean']), 'std': tuple(saved['std'])} |
|
|
| def forward_features(self, x): |
| t = self.vit.forward_features(x)[:, self.vit.num_prefix_tokens:] |
| return t.reshape(t.shape[0], *self.vit.patch_embed.grid_size, -1) |
|
|
| def get_classifier(self): |
| return self.head |
|
|
| def forward(self, x): |
| return self.head(self.forward_features(x).mean(dim=(1, 2))) |
|
|
|
|
| def bhwc_cells(model, x): |
| return model.forward_features(x) |
|
|
|
|
| def clip_cells(model, x): |
| return model.forward_features(x) if isinstance(model, ClipViT) else model.forward_features(x).permute(0, 2, 3, 1) |
|
|
|
|
| def zero_shot_head(model): |
| return model.cell_logits |
|
|
|
|
| MODELS = { |
| |
| 'resnet50': (timm_model('resnet50.tv2_in1k'), cnn_cells, linear_head), |
| 'convnext_small': (timm_model('convnext_small.fb_in1k'), cnn_cells, norm_head), |
| 'swin_small': (timm_model('swin_small_patch4_window7_224.ms_in1k'), swin_cells, linear_head), |
| 'maxvit_tiny': (timm_model('maxvit_tiny_tf_224.in1k'), cnn_cells, norm_tanh_head), |
| 'vit_small': (timm_model('vit_small_patch16_224.augreg_in21k_ft_in1k'), vit_cells, linear_head), |
| 'deit_small': (timm_model('deit_small_patch16_224.fb_in1k'), vit_cells, linear_head), |
| 'deit3_medium': (timm_model('deit3_medium_patch16_224.fb_in1k'), vit_cells, linear_head), |
| |
| 'dino_resnet50': (dino_resnet50, cnn_cells, None), |
| 'dino_vit_small': (timm_model('vit_small_patch16_224.dino'), vit_cells, None), |
| 'dinov2_vit_small': (timm_model('vit_small_patch14_dinov2.lvd142m', img_size=224), vit_cells, None), |
| 'dinov2_reg_vit_small': (timm_model('vit_small_patch14_reg4_dinov2.lvd142m', img_size=224), vit_cells, None), |
| 'dinov3_vit_small_plus': (timm_model('vit_small_plus_patch16_dinov3.lvd1689m', img_size=224), vit_cells, None), |
| 'dinov3_convnext_small': (timm_model('convnext_small.dinov3_lvd1689m'), cnn_cells, None), |
| 'clip_resnet101': (timm_model('resnet101_clip_gap.openai'), cnn_cells, None), |
| 'tinyclip_vit_medium': (timm_model('vit_medium_patch16_clip_224.tinyclip_yfcc15m'), vit_cells, None), |
| 'fcmae_convnextv2_tiny': (timm_model('convnextv2_tiny.fcmae'), cnn_cells, None), |
| |
| 'clip_rn50x4': (ClipRN, clip_cells, zero_shot_head), |
| |
| 'tumour_detector': (tumour_detector, cnn_cells, linear_head), |
| 'kaiko_vit_small': (kaiko_vits16, vit_cells, None), |
| 'ctranspath': (ctranspath, swin_cells, None), |
| 'hibou_b': (HibouB, bhwc_cells, linear_head), |
| 'quiltnet_b32': (ClipViT, clip_cells, zero_shot_head), |
| } |
| SIZES = {'clip_rn50x4': 288} |
| IMAGE_LOGITS = {'clip_rn50x4', 'quiltnet_b32'} |
|
|
|
|
| class DenseReadout(nn.Module): |
| """Normalizes the image and returns the model's feature map and, for a classifier, its dense readout.""" |
|
|
| def __init__(self, model, cells, head, image_logits=False): |
| super().__init__() |
| self.model, self.cells, self.head, self.image_logits = model, cells, head and head(model), image_logits |
| cfg = model.pretrained_cfg |
| self.register_buffer('mean', torch.tensor(cfg['mean'])[:, None, None]) |
| self.register_buffer('std', torch.tensor(cfg['std'])[:, None, None]) |
|
|
| def forward(self, image): |
| cells = self.cells(self.model, (image - self.mean) / self.std) |
| if self.head is None: |
| return cells.permute(0, 3, 1, 2) |
| out = (cells.permute(0, 3, 1, 2), self.head(cells).permute(0, 3, 1, 2)) |
| if self.image_logits: |
| out += (self.model((image - self.mean) / self.std),) |
| return out |
|
|
|
|
| def weights_to_float16(model): |
| """Store large float32 weights as float16, each followed by a Cast back to float32.""" |
| graph = model.graph |
| casts = [] |
| for init in list(graph.initializer): |
| if init.data_type != onnx.TensorProto.FLOAT or np.prod(init.dims) < 1024: |
| continue |
| half = numpy_helper.from_array(numpy_helper.to_array(init).astype(np.float16), init.name + '_fp16') |
| graph.initializer.remove(init) |
| graph.initializer.append(half) |
| casts.append(helper.make_node('Cast', [half.name], [init.name], to=onnx.TensorProto.FLOAT)) |
| nodes = casts + list(graph.node) |
| del graph.node[:] |
| graph.node.extend(nodes) |
| return model |
|
|
|
|
| def export(name): |
| size = SIZES.get(name, 224) |
| create, cells, head = MODELS[name] |
| model = create().eval() |
| dense = DenseReadout(model, cells, head, name in IMAGE_LOGITS).eval() |
| outputs = ['features'] if head is None else ['features', 'cell_logits'] + (['logits'] if name in IMAGE_LOGITS else []) |
| path = OUT / f'{name}.onnx' |
| torch.onnx.export(dense, torch.rand(1, 3, size, size), path, input_names=['image'], |
| output_names=outputs, opset_version=17, dynamo=False) |
| onnx_model = weights_to_float16(onnx.load(path)) |
| onnx.checker.check_model(onnx_model) |
| onnx.save(onnx_model, path) |
|
|
| |
| |
| session = ort.InferenceSession(path) |
| file_error, readout_error, same_top1, corr = 0, 0, 0, 1 |
| images = sorted(SAMPLES.glob('*.jpg')) |
| for p in images: |
| x = torch.from_numpy(np.asarray(Image.open(p).convert('RGB').resize((size, size), Image.BICUBIC))) |
| x = x.permute(2, 0, 1)[None].float() / 255 |
| got_all = session.run(None, {'image': x.numpy()}) |
| got = got_all[min(1, len(got_all) - 1)] |
| with torch.no_grad(): |
| out = dense(x) |
| expected = (out if head is None else out[1]).numpy() |
| file_error = max(file_error, np.abs(got - expected).max() / np.abs(expected).max()) |
| corr = min(corr, np.corrcoef(got.ravel(), expected.ravel())[0, 1]) |
| if head is not None: |
| with torch.no_grad(): |
| logits = model((x - dense.mean) / dense.std).numpy() |
| if name in IMAGE_LOGITS: |
| readout_error = max(readout_error, np.abs(got_all[2] - logits).max()) |
| same_top1 += got_all[2].argmax() == logits.argmax() |
| else: |
| readout_error = max(readout_error, np.abs(got.mean((2, 3)) - logits).max()) |
| same_top1 += got.mean((2, 3)).argmax() == logits.argmax() |
| report = f'{path.name}: {path.stat().st_size / 2**20:.1f} MB, {outputs[min(1, len(outputs) - 1)]} {got.shape[1:]}, ' \ |
| f'max |file - pytorch| / max = {file_error:.4f}, min correlation {corr:.5f}' |
| if head is not None: |
| what = '|file logits - logits|' if name in IMAGE_LOGITS else '|mean(cell_logits) - logits|' |
| report += f', max {what} = {readout_error:.4f}, same top-1 {same_top1}/{len(images)}' |
| print(report, flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| for name in sys.argv[1:] or MODELS: |
| export(name) |
|
|
