Wrapped DINOv2 encoder supporting gradient checkpointing. Input is RGB image in range [0, 1].
| 69 | |
| 70 | |
| 71 | class DINOv2Encoder(nn.Module): |
| 72 | "Wrapped DINOv2 encoder supporting gradient checkpointing. Input is RGB image in range [0, 1]." |
| 73 | backbone: DinoVisionTransformer |
| 74 | image_mean: torch.Tensor |
| 75 | image_std: torch.Tensor |
| 76 | dim_features: int |
| 77 | |
| 78 | def __init__(self, backbone: str, intermediate_layers: Union[int, List[int]], dim_out: int, **deprecated_kwargs): |
| 79 | super(DINOv2Encoder, self).__init__() |
| 80 | |
| 81 | self.intermediate_layers = intermediate_layers |
| 82 | |
| 83 | # Load the backbone |
| 84 | self.hub_loader = getattr(importlib.import_module(".dinov2.hub.backbones", __package__), backbone) |
| 85 | self.backbone_name = backbone |
| 86 | self.backbone = self.hub_loader(pretrained=False) |
| 87 | |
| 88 | self.dim_features = self.backbone.blocks[0].attn.qkv.in_features |
| 89 | self.num_features = intermediate_layers if isinstance(intermediate_layers, int) else len(intermediate_layers) |
| 90 | |
| 91 | self.output_projections = nn.ModuleList([ |
| 92 | nn.Conv2d(in_channels=self.dim_features, out_channels=dim_out, kernel_size=1, stride=1, padding=0,) |
| 93 | for _ in range(self.num_features) |
| 94 | ]) |
| 95 | |
| 96 | self.register_buffer("image_mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)) |
| 97 | self.register_buffer("image_std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)) |
| 98 | |
| 99 | def init_weights(self): |
| 100 | pretrained_backbone_state_dict = self.hub_loader(pretrained=True).state_dict() |
| 101 | self.backbone.load_state_dict(pretrained_backbone_state_dict) |
| 102 | |
| 103 | def enable_gradient_checkpointing(self): |
| 104 | for i in range(len(self.backbone.blocks)): |
| 105 | wrap_module_with_gradient_checkpointing(self.backbone.blocks[i]) |
| 106 | |
| 107 | def enable_pytorch_native_sdpa(self): |
| 108 | for i in range(len(self.backbone.blocks)): |
| 109 | wrap_dinov2_attention_with_sdpa(self.backbone.blocks[i].attn) |
| 110 | |
| 111 | def forward(self, image: torch.Tensor, token_rows: int, token_cols: int, return_class_token: bool = False) -> Tuple[torch.Tensor, torch.Tensor]: |
| 112 | image_14 = F.interpolate(image, (token_rows * 14, token_cols * 14), mode="bilinear", align_corners=False, antialias=True) |
| 113 | image_14 = (image_14 - self.image_mean) / self.image_std |
| 114 | |
| 115 | # Get intermediate layers from the backbone |
| 116 | features = self.backbone.get_intermediate_layers(image_14, n=self.intermediate_layers, return_class_token=True) |
| 117 | |
| 118 | # Project features to the desired dimensionality |
| 119 | x = torch.stack([ |
| 120 | proj(feat.permute(0, 2, 1).unflatten(2, (token_rows, token_cols)).contiguous()) |
| 121 | for proj, (feat, clstoken) in zip(self.output_projections, features) |
| 122 | ], dim=1).sum(dim=1) |
| 123 | |
| 124 | if return_class_token: |
| 125 | return x, features[-1][1] |
| 126 | else: |
| 127 | return x |
| 128 |