CelebAMask-HQ的SOTA实现：高精度人脸分割的工程优化实践

1次阅读

共计 2115 个字符，预计需要花费 6 分钟才能阅读完成。

CelebAMask-HQ 作为目前最精细的人脸分割数据集，包含 30,000 张高分辨率图像（1024×1024）和 19 类语义标签。但在实际业务落地时面临三个核心问题：

细粒度分割精度要求：发丝、饰品等微小区域的分割直接影响用户体验，传统模型在 hair、earrings 等类别 IOU 普遍低于 85%
高分辨率内存压力：单张图像全分辨率加载需要 4GB+ 显存，批量训练时 OOM 问题频发
实时性瓶颈：1080P 视频处理需要达到 25FPS，原模型在 3090 上仅 8FPS

对比主流分割模型在 CelebAMask-HQ 验证集的表现：

模型	mIOU(%)	参数量(M)	推理速度(ms)
U-Net++	89.7	24.1	68
DeepLabv3+	92.3	15.8	53
HRNet	91.5	48.6	112

最终选择 DeepLabv3+ 为基础架构，主要改进点：

将原版 ASPP 替换为可变形卷积 DCNv2，提升对不规则边缘的建模能力
在 decoder 阶段添加 CBAM 注意力模块，增强头发等细小区域的特征提取
采用混合空洞卷积策略（HDC）缓解网格伪影

class CBAM_ASPP(nn.Module):
    def __init__(self, in_channels, dilation_rates=[6,12,18]):
        super().__init__()
        # 通道注意力
        self.channel_att = nn.Sequential(nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, in_channels//8, 1),
            nn.ReLU(),
            nn.Conv2d(in_channels//8, in_channels, 1),
            nn.Sigmoid())
        # 空间注意力
        self.spatial_att = nn.Sequential(nn.Conv2d(2, 1, 7, padding=3),
            nn.Sigmoid())
        # 空洞卷积分支
        self.conv_branches = nn.ModuleList([
            nn.Conv2d(in_channels, in_channels, 3, 
                      padding=d, dilation=d) for d in dilation_rates
        ])

    def forward(self, x):
        # 通道注意力计算
        channel_att = self.channel_att(x)
        # 空间注意力计算
        max_pool = torch.max(x, dim=1, keepdim=True)[0]
        mean_pool = torch.mean(x, dim=1, keepdim=True)
        spatial_att = self.spatial_att(torch.cat([max_pool, mean_pool], dim=1))

        # 多尺度特征融合
        branch_outputs = [conv(x) for conv in self.conv_branches]
        features = torch.cat([x] + branch_outputs, dim=1)

        # 双重注意力加权
        return features * channel_att * spatial_att

数据增强策略：
随机擦除 (RandomErasing) 针对遮挡场景
ElasticTransform 增强发丝边缘
颜色抖动限制在 ΔE<5 避免影响语义

损失函数组合：

class HybridLoss(nn.Module):
    def __init__(self, class_weights):
        super().__init__()
        self.dice = DiceLoss()
        self.focal = FocalLoss(weight=class_weights)

    def forward(self, pred, target):
        return 0.4*self.dice(pred, target) + 0.6*self.focal(pred, target)