| 299 | |
| 300 | |
| 301 | class VisionTransformerWithLinear(nn.Module): |
| 302 | |
| 303 | def __init__(self, base_vit, num_classes=200): |
| 304 | |
| 305 | super().__init__() |
| 306 | |
| 307 | self.base_vit = base_vit |
| 308 | self.fc = nn.Linear(768, num_classes) |
| 309 | |
| 310 | def forward(self, x, return_features=False): |
| 311 | |
| 312 | features = self.base_vit(x) |
| 313 | features = torch.nn.functional.normalize(features, dim=-1) |
| 314 | logits = self.fc(features) |
| 315 | |
| 316 | if return_features: |
| 317 | return logits, features |
| 318 | else: |
| 319 | return logits |
| 320 | |
| 321 | @torch.no_grad() |
| 322 | def normalize_prototypes(self): |
| 323 | w = self.fc.weight.data.clone() |
| 324 | w = torch.nn.functional.normalize(w, dim=1, p=2) |
| 325 | self.fc.weight.copy_(w) |
| 326 | |
| 327 | |
| 328 |
nothing calls this directly
no outgoing calls
no test coverage detected