Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| from .encoder import Encoder | |
| class ViTClassifier(nn.Module): | |
| """ | |
| Vision Transformer for Classification (Fine-tuning MAE). | |
| """ | |
| def __init__( | |
| self, | |
| encoder: Encoder, | |
| num_classes: int = 10, | |
| global_pool: bool = False | |
| ): | |
| super().__init__() | |
| self.encoder = encoder | |
| self.num_classes = num_classes | |
| self.global_pool = global_pool | |
| # Use encoder's embed_dim for the head | |
| embed_dim = encoder.cls_token.shape[-1] | |
| if global_pool: | |
| self.fc_norm = nn.LayerNorm(embed_dim) | |
| else: | |
| self.fc_norm = None | |
| self.head = nn.Linear(embed_dim, num_classes) | |
| # Initialize head | |
| torch.nn.init.xavier_uniform_(self.head.weight) | |
| nn.init.constant_(self.head.bias, 0) | |
| def forward_features(self, x): | |
| x = self.encoder.patch_embed(x) | |
| x = x + self.encoder.pos_embed[:, 1:, :] | |
| cls_token = self.encoder.cls_token + self.encoder.pos_embed[:, :1, :] | |
| cls_tokens = cls_token.expand(x.shape[0], -1, -1) | |
| x = torch.cat((cls_tokens, x), dim=1) | |
| for blk in self.encoder.blocks: | |
| x = blk(x) | |
| x = self.encoder.norm(x) | |
| return x[:, 0] # CLS embedding | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| # Encoder forward (no masking during fine-tuning) | |
| # x, _, _ = self.encoder(x, mask_ratio=0.0) | |
| # Manually call encoder parts to avoid masking logic overhead if needed | |
| # but for simplicity, we use the forward method | |
| latent, _, _ = self.encoder(x, mask_ratio=0.0) | |
| if self.global_pool: | |
| # Global average pooling over all tokens (excluding cls token) | |
| x = latent[:, 1:, :].mean(dim=1) | |
| x = self.fc_norm(x) | |
| else: | |
| # Use only the cls token | |
| x = latent[:, 0] | |
| x = self.head(x) | |
| return x | |