| """Backbone builders for scratch face-recognition experiments.""" |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Sequence |
|
|
| import torch |
| from facenet_pytorch import InceptionResnetV1 |
| from torch import nn |
|
|
|
|
| class Flatten(nn.Module): |
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return x.view(x.size(0), -1) |
|
|
|
|
| class IRBlock(nn.Module): |
| def __init__(self, in_channels: int, out_channels: int, stride: int): |
| super().__init__() |
| self.residual = nn.Sequential( |
| nn.BatchNorm2d(in_channels), |
| nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False), |
| nn.BatchNorm2d(out_channels), |
| nn.PReLU(out_channels), |
| nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False), |
| nn.BatchNorm2d(out_channels), |
| ) |
| if in_channels == out_channels and stride == 1: |
| self.shortcut = nn.Identity() |
| else: |
| self.shortcut = nn.Sequential( |
| nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False), |
| nn.BatchNorm2d(out_channels), |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.residual(x) + self.shortcut(x) |
|
|
|
|
| class IRResNet(nn.Module): |
| """InsightFace-style IR-ResNet for 112x112 face crops.""" |
|
|
| def __init__(self, layers: Sequence[int], embedding_size: int = 512, dropout: float = 0.4): |
| super().__init__() |
| self.input_layer = nn.Sequential( |
| nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False), |
| nn.BatchNorm2d(64), |
| nn.PReLU(64), |
| ) |
| self.body = nn.Sequential( |
| self._make_stage(64, 64, layers[0], stride=2), |
| self._make_stage(64, 128, layers[1], stride=2), |
| self._make_stage(128, 256, layers[2], stride=2), |
| self._make_stage(256, 512, layers[3], stride=2), |
| ) |
| self.output_layer = nn.Sequential( |
| nn.BatchNorm2d(512), |
| nn.Dropout(dropout), |
| Flatten(), |
| nn.Linear(512 * 7 * 7, embedding_size), |
| nn.BatchNorm1d(embedding_size), |
| ) |
| self._init_weights() |
|
|
| def _make_stage(self, in_channels: int, out_channels: int, blocks: int, stride: int) -> nn.Sequential: |
| layers = [IRBlock(in_channels, out_channels, stride)] |
| for _ in range(1, blocks): |
| layers.append(IRBlock(out_channels, out_channels, 1)) |
| return nn.Sequential(*layers) |
|
|
| def _init_weights(self) -> None: |
| for module in self.modules(): |
| if isinstance(module, nn.Conv2d): |
| nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu") |
| elif isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d)): |
| nn.init.ones_(module.weight) |
| nn.init.zeros_(module.bias) |
| elif isinstance(module, nn.Linear): |
| nn.init.xavier_uniform_(module.weight) |
| nn.init.zeros_(module.bias) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = self.input_layer(x) |
| x = self.body(x) |
| return self.output_layer(x) |
|
|
|
|
| def canonical_backbone_name(name: str) -> str: |
| normalized = name.lower().replace("-", "_") |
| aliases = { |
| "inceptionresnetv1": "inception_resnet_v1", |
| "inception_resnet": "inception_resnet_v1", |
| "inception_resnet_v1": "inception_resnet_v1", |
| "ir18": "ir_resnet18", |
| "ir_resnet18": "ir_resnet18", |
| "ir34": "ir_resnet34", |
| "ir_resnet34": "ir_resnet34", |
| } |
| if normalized not in aliases: |
| raise ValueError(f"Unsupported backbone: {name}") |
| return aliases[normalized] |
|
|
|
|
| def build_backbone(name: str, pretrained_model: str | None = None) -> tuple[nn.Module, int, str]: |
| canonical = canonical_backbone_name(name) |
| if canonical == "inception_resnet_v1": |
| model = InceptionResnetV1(pretrained=pretrained_model, classify=False) |
| elif canonical == "ir_resnet18": |
| model = IRResNet([2, 2, 2, 2]) |
| elif canonical == "ir_resnet34": |
| model = IRResNet([3, 4, 6, 3]) |
| else: |
| raise ValueError(f"Unsupported backbone: {name}") |
| return model, 512, canonical |
|
|
