| import torch.nn as nn | |
| from ...models.dino.head import DINOHead | |
| class iBOTHead(DINOHead): | |
| def __init__( | |
| self, | |
| *args, | |
| patch_out_dim=8192, | |
| n_layers=3, | |
| hidden_dim=2048, | |
| bottleneck_dim=256, | |
| norm_last_layer=True, | |
| shared_head=False, | |
| **kwargs | |
| ): | |
| super(iBOTHead, self).__init__( | |
| *args, | |
| n_layers=n_layers, | |
| hidden_dim=hidden_dim, | |
| bottleneck_dim=bottleneck_dim, | |
| norm_last_layer=norm_last_layer, | |
| **kwargs | |
| ) | |
| if not shared_head: | |
| if bottleneck_dim > 0: | |
| self.last_layer2 = nn.utils.weight_norm( | |
| nn.Linear(bottleneck_dim, patch_out_dim, bias=False) | |
| ) | |
| self.last_layer2.weight_g.data.fill_(1) | |
| if norm_last_layer: | |
| self.last_layer2.weight_g.requires_grad = False | |
| else: | |
| self.mlp2 = nn.Linear(hidden_dim, patch_out_dim) | |
| self.last_layer2 = None | |
| else: | |
| if bottleneck_dim > 0: | |
| self.last_layer2 = self.last_layer | |
| else: | |
| self.mlp2 = self.mlp[-1] | |
| self.last_layer2 = None | |
| def forward(self, x): | |
| if len(x.shape) == 2: | |
| return super(iBOTHead, self).forward(x) | |
| if self.last_layer is not None: | |
| x = self.mlp(x) | |
| x = nn.functional.normalize(x, dim=-1, p=2) | |
| x1 = self.last_layer(x[:, 0]) | |
| x2 = self.last_layer2(x[:, 1:]) | |
| else: | |
| x = self.mlp[:-1](x) | |
| x1 = self.mlp[-1](x[:, 0]) | |
| x2 = self.mlp2(x[:, 1:]) | |
| return x1, x2 | |
| def _build_norm(self, norm, hidden_dim, **kwargs): | |
| if norm == "bn": | |
| norm = nn.BatchNorm1d(hidden_dim, **kwargs) | |
| elif norm == "syncbn": | |
| norm = nn.SyncBatchNorm(hidden_dim, **kwargs) | |
| elif norm == "ln": | |
| norm = nn.LayerNorm(hidden_dim, **kwargs) | |
| else: | |
| assert norm is None, "unknown norm type {}".format(norm) | |
| return norm | |
| def _build_act(self, act): | |
| if act == "relu": | |
| act = nn.ReLU() | |
| elif act == "gelu": | |
| act = nn.GELU() | |
| else: | |
| assert False, "unknown act type {}".format(act) | |
| return act | |