VisionLanguageGroup/MicroscopyMatching

1"""2Copyright © 2025 Howard Hughes Medical Institute, Authored by Carsen Stringer and Marius Pachitariu.3"""4 5import torch6from segment_anything import sam_model_registry7torch.backends.cuda.matmul.allow_tf32 = True8from torch import nn 9import torch.nn.functional as F10 11class Transformer(nn.Module):12    def __init__(self, backbone="vit_l", ps=8, nout=3, bsize=256, rdrop=0.4,13                  checkpoint=None, dtype=torch.float32):14        super(Transformer, self).__init__()15        """16        print(self.encoder.patch_embed)17            PatchEmbed(18            (proj): Conv2d(3, 1024, kernel_size=(16, 16), stride=(16, 16))19            )20        print(self.encoder.neck)21            Sequential(22            (0): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)23            (1): LayerNorm2d()24            (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)25            (3): LayerNorm2d()26            )27        """28        # instantiate the vit model, default to not loading SAM29        # checkpoint = sam_vit_l_0b3195.pth is standard pretrained SAM30        self.encoder = sam_model_registry[backbone](checkpoint).image_encoder31        w = self.encoder.patch_embed.proj.weight.detach()32        nchan = w.shape[0]33        34        # change token size to ps x ps35        self.ps = ps36        self.encoder.patch_embed.proj = nn.Conv2d(3, nchan, stride=ps, kernel_size=ps)37        self.encoder.patch_embed.proj.weight.data = w[:,:,::16//ps,::16//ps]38        39        # adjust position embeddings for new bsize and new token size40        ds = (1024 // 16) // (bsize // ps)41        self.encoder.pos_embed = nn.Parameter(self.encoder.pos_embed[:,::ds,::ds], requires_grad=True)42 43        # readout weights for nout output channels44        # if nout is changed, weights will not load correctly from pretrained Cellpose-SAM45        self.nout = nout46        self.out = nn.Conv2d(256, self.nout * ps**2, kernel_size=1)47 48        # W2 reshapes token space to pixel space, not trainable49        self.W2 = nn.Parameter(torch.eye(self.nout * ps**2).reshape(self.nout*ps**2, self.nout, ps, ps), 50                               requires_grad=False)51        52        # fraction of layers to drop at random during training53        self.rdrop = rdrop54 55        # average diameter of ROIs from training images from fine-tuning 56        self.diam_labels = nn.Parameter(torch.tensor([30.]), requires_grad=False)57        # average diameter of ROIs during main training58        self.diam_mean = nn.Parameter(torch.tensor([30.]), requires_grad=False)59        60        # set attention to global in every layer61        for blk in self.encoder.blocks:62            blk.window_size = 063 64        self.dtype = dtype65 66    def forward(self, x, feat=None):      67        # same progression as SAM until readout68        x = self.encoder.patch_embed(x)69        if feat is not None:70            feat = self.encoder.patch_embed(feat)71            x = x + x * feat * 0.572        73        if self.encoder.pos_embed is not None:74            x = x + self.encoder.pos_embed75        76        if self.training and self.rdrop > 0:77            nlay = len(self.encoder.blocks)78            rdrop = (torch.rand((len(x), nlay), device=x.device) < 79                     torch.linspace(0, self.rdrop, nlay, device=x.device)).to(x.dtype)80            for i, blk in enumerate(self.encoder.blocks):            81                mask = rdrop[:,i].unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)82                x = x * mask + blk(x) * (1-mask)83        else:84            for blk in self.encoder.blocks:85                x = blk(x)86 87        x = self.encoder.neck(x.permute(0, 3, 1, 2))88 89        # readout is changed here90        x1 = self.out(x)91        x1 = F.conv_transpose2d(x1, self.W2, stride = self.ps, padding = 0)92        93        # maintain the second output of feature size 256 for backwards compatibility94           95        return x1, torch.randn((x.shape[0], 256), device=x.device)96    97    def load_model(self, PATH, device, strict = False):        98        state_dict = torch.load(PATH, map_location = device, weights_only=True)99        keys = [k for k in state_dict.keys()]100        if keys[0][:7] == "module.":101            from collections import OrderedDict102            new_state_dict = OrderedDict()103            for k, v in state_dict.items():104                name = k[7:] # remove 'module.' of DataParallel/DistributedDataParallel105                new_state_dict[name] = v106            self.load_state_dict(new_state_dict, strict = strict)107        else:108            self.load_state_dict(state_dict, strict = strict)109 110        if self.dtype != torch.float32:111            self = self.to(self.dtype)112 113    114    @property115    def device(self):116        """117        Get the device of the model.118 119        Returns:120            torch.device: The device of the model.121        """122        return next(self.parameters()).device123 124    def save_model(self, filename):125        """126        Save the model to a file.127 128        Args:129            filename (str): The path to the file where the model will be saved.130        """131        torch.save(self.state_dict(), filename)132 133