| 83 | |
| 84 | class DINOHead(nn.Module): |
| 85 | def __init__(self, in_dim, out_dim, use_bn=False, norm_last_layer=True, nlayers=3, hidden_dim=2048, bottleneck_dim=256): |
| 86 | super().__init__() |
| 87 | nlayers = max(nlayers, 1) |
| 88 | if nlayers == 1: |
| 89 | self.mlp = nn.Linear(in_dim, bottleneck_dim) |
| 90 | else: |
| 91 | layers = [nn.Linear(in_dim, hidden_dim)] |
| 92 | if use_bn: |
| 93 | layers.append(nn.BatchNorm1d(hidden_dim)) |
| 94 | layers.append(nn.GELU()) |
| 95 | for _ in range(nlayers - 2): |
| 96 | layers.append(nn.Linear(hidden_dim, hidden_dim)) |
| 97 | if use_bn: |
| 98 | layers.append(nn.BatchNorm1d(hidden_dim)) |
| 99 | layers.append(nn.GELU()) |
| 100 | layers.append(nn.Linear(hidden_dim, bottleneck_dim)) |
| 101 | self.mlp = nn.Sequential(*layers) |
| 102 | self.apply(self._init_weights) |
| 103 | self.last_layer = nn.utils.weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False)) |
| 104 | self.last_layer.weight_g.data.fill_(1) |
| 105 | if norm_last_layer: |
| 106 | self.last_layer.weight_g.requires_grad = False |
| 107 | |
| 108 | def _init_weights(self, m): |
| 109 | if isinstance(m, nn.Linear): |