Self-attention block with FFN for DINOv3.
| 135 | |
| 136 | |
| 137 | class SelfAttentionBlock(nn.Module): |
| 138 | """Self-attention block with FFN for DINOv3.""" |
| 139 | |
| 140 | def __init__( |
| 141 | self, |
| 142 | dim: int, |
| 143 | num_heads: int, |
| 144 | ffn_ratio: float = 4.0, |
| 145 | qkv_bias: bool = False, |
| 146 | proj_bias: bool = True, |
| 147 | ffn_bias: bool = True, |
| 148 | drop: float = 0.0, |
| 149 | attn_drop: float = 0.0, |
| 150 | init_values=None, |
| 151 | drop_path: float = 0.0, |
| 152 | act_layer: Callable[..., nn.Module] = nn.GELU, |
| 153 | norm_layer: Callable[..., nn.Module] = nn.LayerNorm, |
| 154 | attn_class: Callable[..., nn.Module] = SelfAttention, |
| 155 | ffn_layer: Callable[..., nn.Module] = Mlp, |
| 156 | mask_k_bias: bool = False, |
| 157 | device=None, |
| 158 | use_qk_norm: bool = False, |
| 159 | ) -> None: |
| 160 | super().__init__() |
| 161 | self.norm1 = norm_layer(dim) |
| 162 | self.attn = attn_class( |
| 163 | dim, |
| 164 | num_heads=num_heads, |
| 165 | qkv_bias=qkv_bias, |
| 166 | proj_bias=proj_bias, |
| 167 | attn_drop=attn_drop, |
| 168 | proj_drop=drop, |
| 169 | mask_k_bias=mask_k_bias, |
| 170 | device=device, |
| 171 | use_qk_norm=use_qk_norm, |
| 172 | ) |
| 173 | self.ls1 = LayerScale(dim, init_values=init_values, device=device) if init_values else nn.Identity() |
| 174 | |
| 175 | self.norm2 = norm_layer(dim) |
| 176 | mlp_hidden_dim = int(dim * ffn_ratio) |
| 177 | self.mlp = ffn_layer( |
| 178 | in_features=dim, |
| 179 | hidden_features=mlp_hidden_dim, |
| 180 | act_layer=act_layer, |
| 181 | drop=drop, |
| 182 | bias=ffn_bias, |
| 183 | device=device, |
| 184 | ) |
| 185 | self.ls2 = LayerScale(dim, init_values=init_values, device=device) if init_values else nn.Identity() |
| 186 | |
| 187 | self.sample_drop_ratio = drop_path |
| 188 | |
| 189 | @staticmethod |
| 190 | def _maybe_index_rope(rope: Optional[Tuple[Tensor, Tensor]], indices: Tensor) -> Optional[Tuple[Tensor, Tensor]]: |
| 191 | if rope is None: |
| 192 | return None |
| 193 | |
| 194 | sin, cos = rope |