File size: 2,563 Bytes
db208b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """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:
# The stride-16 stage: 448 px input gives a 28x28 map, 7x7 per cell.
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),
# Max, not mean: any part of the target selects the cell.
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)
|