File size: 4,296 Bytes
11d0468 | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | """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
|