👥 Age + Gender — v20260507 Final（8hr 自主研究）

import torch, torch.nn as nn, timm
from PIL import Image
import torchvision.transforms as T

class MultiHead(nn.Module):
    def __init__(self, backbone_name, drop_rate=0.3, num_age=4):
        super().__init__()
        self.backbone = timm.create_model(backbone_name, pretrained=False,
                                          num_classes=0, global_pool="avg",
                                          drop_rate=drop_rate)
        with torch.no_grad():
            feat_dim = self.backbone(torch.zeros(1,3,64,64)).shape[-1]
        self.gender_head = nn.Linear(feat_dim, 1)
        self.age_head = nn.Linear(feat_dim, num_age)
    def forward(self, x):
        f = self.backbone(x)
        return self.gender_head(f).squeeze(-1), self.age_head(f)

ckpt = torch.load("age_gender_v20260507_swa4.pt", weights_only=False)
model = MultiHead(ckpt["args"]["backbone"]).eval()
model.load_state_dict(ckpt["model_state"])

mean = [0.485, 0.456, 0.406]; std = [0.229, 0.224, 0.225]
tf = T.Compose([T.Resize((384, 192)), T.ToTensor(), T.Normalize(mean, std)])
x = tf(Image.open("person_crop.jpg").convert("RGB")).unsqueeze(0)
with torch.no_grad():
    g_logit, a_logit = model(x)
gender = "female" if torch.sigmoid(g_logit).item() > 0.5 else "male"
age = ["child","young","adult","elder"][a_logit.argmax(dim=-1).item()]
# 注意：young 在 PA-100K + MSP60K 都沒 supervision，推論不會出 young