Merge pull request #4 from recursionpharma/more-code

.vscode/settings.json

@@ -0,0 +1,5 @@
+{
+    "flake8.args": [
+        "--max-line-length=120"
+    ]
+}

requirements.txt

@@ -0,0 +1,9 @@
+huggingface-hub==0.18.0
+timm==0.9.7
+torch==2.1.0+cu121
+torchmetrics==1.2.0
+torchvision==0.16.0+cu121
+tqdm==4.66.1
+transformers==4.35.2
+xformers==0.0.22.post7
+zarr==2.16.1

vit_encoder.py

@@ -0,0 +1,60 @@
+from typing import Dict
+
+import timm.models.vision_transformer as vit
+import torch
+
+
+def build_imagenet_baselines() -> Dict[str, torch.jit.ScriptModule]:
+    """This returns the prepped imagenet encoders from timm, not bad for microscopy data."""
+    vit_backbones = [
+        _make_vit(vit.vit_small_patch16_384),
+        _make_vit(vit.vit_base_patch16_384),
+        _make_vit(vit.vit_base_patch8_224),
+        _make_vit(vit.vit_large_patch16_384),
+    ]
+    model_names = [
+        "vit_small_patch16_384",
+        "vit_base_patch16_384",
+        "vit_base_patch8_224",
+        "vit_large_patch16_384",
+    ]
+    imagenet_encoders = list(map(_make_torchscripted_encoder, vit_backbones))
+    return {name: model for name, model in zip(model_names, imagenet_encoders)}
+
+
+def _make_torchscripted_encoder(vit_backbone) -> torch.jit.ScriptModule:
+    dummy_input = torch.testing.make_tensor(
+        (2, 6, 256, 256),
+        low=0,
+        high=255,
+        dtype=torch.uint8,
+        device=torch.device("cpu"),
+    )
+    encoder = torch.nn.Sequential(
+        Normalizer(),
+        torch.nn.LazyInstanceNorm2d(
+            affine=False, track_running_stats=False
+        ),  # this module performs self-standardization, very important
+        vit_backbone,
+    ).to(device="cpu")
+    _ = encoder(dummy_input)  # get those lazy modules built
+    return torch.jit.freeze(torch.jit.script(encoder.eval()))
+
+
+def _make_vit(constructor):
+    return constructor(
+        pretrained=True,  # download imagenet weights
+        img_size=256,  # 256x256 crops
+        in_chans=6,  # we expect 6-channel microscopy images
+        num_classes=0,
+        fc_norm=None,
+        class_token=True,
+        global_pool="avg",  # minimal perf diff btwn "cls" and "avg"
+    )
+
+
+class Normalizer(torch.nn.Module):
+    def forward(self, pixels: torch.Tensor) -> torch.Tensor:
+        pixels = pixels.float()
+        pixels /= 255.0
+        return pixels