Class SAMWrapper

examples/common/model_loader.py:198–227

class SAMWrapper(nn.Module):
    def __init__(self, sam_encoder: nn.Module):
        super().__init__()
        self.inner = sam_encoder

    @property
    def embed_dim(self):
        return self.inner.patch_embed.proj.out_channels

    @property
    def patch_size(self):
        return self.inner.patch_embed.proj.kernel_size[0]

    @property
    def vision_encoder(self):
        return self.inner

    def forward(self, x: torch.Tensor, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
        x = self.inner.patch_embed(x)
        if self.inner.pos_embed is not None:
            x = x + self.inner.pos_embed

        for blk in self.inner.blocks:
            x = blk(x)

        features = x.flatten(1, 2)

        summary = features.mean(dim=1)

        return summary, features


class InternViTWrapper(nn.Module):