Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| from torchvision import models | |
| class CNNEncoder(nn.Module): | |
| def __init__(self, encoded_dim=2048, d_model=512, freeze_backbone=True): | |
| super().__init__() | |
| resnet = models.resnet50(pretrained=True) | |
| self.backbone = nn.Sequential(*list(resnet.children())[:-2]) | |
| self.proj = nn.Linear(encoded_dim, d_model) | |
| self.freeze_backbone = freeze_backbone | |
| if freeze_backbone: | |
| for p in self.backbone.parameters(): | |
| p.requires_grad = False | |
| def unfreeze_backbone(self, unfreeze=True): | |
| self.freeze_backbone = not unfreeze | |
| for p in self.backbone.parameters(): | |
| p.requires_grad = unfreeze | |
| print(f"ResNet Backbone {'Unfrozen' if unfreeze else 'Frozen'}") | |
| def forward(self, images): | |
| """ | |
| images: (B, 3, 224, 224) | |
| return: (B, 49, d_model) | |
| """ | |
| features = self.backbone(images) | |
| # (B, 2048, 7, 7) | |
| B, C, H, W = features.shape | |
| features = features.view(B, C, H * W).permute(0, 2, 1) | |
| # (B, 49, 2048) | |
| features = self.proj(features) | |
| # (B, 49, d_model) | |
| return features |