| """Small pretrained models for independent tiles and continuous grids.""" |
|
|
| import timm |
| from torch import nn |
|
|
| BACKBONE = "mobilenetv4_conv_medium.e250_r384_in12k_ft_in1k" |
|
|
|
|
| class GridModel(nn.Module): |
| def __init__(self, tiles, classes, pretrained=True, backbone=BACKBONE): |
| super().__init__() |
| if tiles not in (9, 16): |
| raise ValueError("Grid size must be 9 or 16 tiles") |
| self.tiles = tiles |
| if tiles == 9: |
| self.backbone = timm.create_model(backbone, pretrained=pretrained) |
| self.backbone.reset_classifier(0) |
| channels = self.backbone.head_hidden_size |
| self.head = nn.Linear(channels, classes) |
| else: |
| |
| stage = ( |
| timm.create_model(backbone, features_only=True).feature_info.reduction().index(16) |
| ) |
| self.backbone = timm.create_model( |
| backbone, pretrained=pretrained, features_only=True, out_indices=(stage,) |
| ) |
| channels = self.backbone.feature_info.channels()[0] |
| self.head = nn.Sequential( |
| nn.Conv2d(channels, 128, 3, padding=1), |
| nn.ReLU(), |
| nn.Conv2d(128, classes, 1), |
| |
| nn.AdaptiveMaxPool2d((4, 4)), |
| ) |
|
|
| def forward(self, images): |
| features = self.backbone(images) |
| return self.head(features if self.tiles == 9 else features[0]) |
|
|
| def freeze_backbone(self, frozen): |
| for parameter in self.backbone.parameters(): |
| parameter.requires_grad_(not frozen) |
|
|
|
|
| def select_target(logits, targets): |
| """Only the requested target is supervised; other targets are unknown.""" |
| import torch |
|
|
| return logits[torch.arange(logits.shape[0], device=logits.device), targets] |
|
|
|
|
| def transfer_backbone(model, state): |
| """Copy matching backbone weights, e.g. from a tile classifier into the 4x4 feature extractor.""" |
| source = {k.removeprefix("backbone."): v for k, v in state.items() if k.startswith("backbone.")} |
| target = model.backbone.state_dict() |
| matched = {k: v for k, v in source.items() if k in target and target[k].shape == v.shape} |
| missing = sorted(set(target) - set(matched)) |
| if missing: |
| raise ValueError( |
| f"Backbone weights missing for {len(missing)} parameters, e.g. {missing[:3]}" |
| ) |
| model.backbone.load_state_dict(matched) |
| return len(matched) |
|
|
