一种人体属性识别的网络结构

提供了一种人体属性识别的网络设计思路，并绘制了网络结构图，给出了实现代码。

目录

0、前言

1、网络结构

2、代码

---

0、前言

        人体属性识别，是一个典型的多标签分类场景。每个人体有多个标签，如年龄、性别、衣着颜色等，而每个属性又有多种类别，如年龄分儿童青年老人、性别分男女、颜色分红绿青蓝紫...

本文提供了一个网络结构来执行这种任务。

1、网络结构

        如图所示，网络接受一个批量的图像，经过backbone（骨干网络，可以是当前已开源的各种网络，如ResNet、EfficientNet、RepVGG等）提取特征，然后将此特征送入后续多个heads，每个head代表一个属性（性别、年龄等）的分类器（使用线性层构建），最后每个head即可输出在该属性上各个类别的概率（如性别head输出[女性概率，男性概率]）：

支持的属性集以及每个属性对应的类别如下：

ATTRIBUTE_ALL = [
    'Gender',
    'Age',
    'UpColor',
    'DownColor',
    'HairLength',
    'Hat',
    'Glasses',
    'SleevesLength',
    'DownLength',
    'Shoes',
    'Bag',
    'BackPack',
    'HandBag',
    'ShoulderBag',
]

CLASSES_ALL = [
    ['Female', 'Male'],
    ['Child', 'Young', 'Adult', 'Old'],
    ['Unknown', 'UpBlack', 'UpWhite', 'UpRed', 'UpPurple', 'UpYellow', 'UpGray', 'UpBlue', 'UpGreen', 'UpBrown'],
    ['Unknown', 'DownBlack', 'DownWhite', 'DownRed', 'DownPurple', 'DownYellow', 'DownGray', 'DownBlue', 'DownGreen', 'DownBrown'],
    ['ShortHair', 'LongHair'],
    ['WithoutHat', 'WithHat'],
    ['WithoutGlasses', 'WithGlasses'],
    ['ShortSleeves', 'LongSleeves'],
    ['DownShort', 'DownLong'],
    ['WithoutShoes', 'WithShoes'],
    ['WithoutBag', 'WithBag'],
    ['WithoutBackPack', 'WithBackPack'],
    ['WithoutHandBag', 'WithHandBag'],
    ['WithoutShoulderBag', 'WithShoulderBag'],
]

2、代码

        我们支持多种backbone，且支持自定义heads，这样可以构建自由的backbone+heads组合：

import os
import sys

sys.path.insert(0, os.getcwd())

import torch

from models.heads.classifier import BaseClassifier
from models.backbones.repvgg import get_RepVGG_func_by_name
from models.backbones.resnet import resnet18, resnet34, resnet50, resnet101
from models.base import FeatClassifier
from data.unify_attributes import ATTRIBUTE_ALL, CLASSES_ALL

BACKBONE_dict = {
    "ResNet18": resnet18,
    "ResNet34": resnet34,
    "ResNet50": resnet50,
    "ResNet101": resnet101,

    "RepVGG": get_RepVGG_func_by_name,
}

RepVGG_pretrained_models = {
    "RepVGG-A0": "RepVGG-A0-train.pth",
    "RepVGG-A1": "RepVGG-A1-train.pth",
    "RepVGG-A2": "RepVGG-A2-train.pth",
    "RepVGG-B2g4": "RepVGG-B2g4-200epochs-train.pth"
}

HEAD_classesNum = {
    "Gender": 2,
    "Age": 4,
    "UpColor": 10,
    "DownColor": 10,
    "HairLength": 2,
    "Hat": 2,
    "Glasses": 2,
    "SleevesLength": 2,
    "DownLength": 2,
    "Shoes": 2,
    "Bag": 2,
    "BackPack": 2,
    "HandBag": 2,
    "ShoulderBag": 2,
}


def build_network(backbone_name, head_names, phase='train', pretrained=False, pretrained_ckpt=None):
    """
    build network by backbone and head names.
    :param backbone_name: str
    :param head_names: list
    :param phase: str, 'train' or other(means 'test')
    :return:
    """
    # build backbone
    if 'repvgg' in backbone_name.lower():
        deploy = False if phase == 'train' else True
        backbone = BACKBONE_dict['RepVGG'](backbone_name)(deploy=deploy)
        if pretrained and not deploy:
            pretrained_ckpt = os.path.join(pretrained_ckpt, RepVGG_pretrained_models[backbone_name])
            state_dict = torch.load(pretrained_ckpt)
            backbone.load_state_dict(state_dict, strict=False)
    elif 'resnet' in backbone_name.lower():
        backbone = BACKBONE_dict[backbone_name](pretrained=pretrained)
    else:
        assert False, "Not supported architecture: {}!".format(backbone_name)

    # build heads
    heads = []
    for h in head_names:
        heads.append(
            BaseClassifier(
                num_classes=len(CLASSES_ALL[ATTRIBUTE_ALL.index(h)]),
                in_planes=backbone.in_planes,
                name=h)
        )

    return FeatClassifier(backbone, heads)


if __name__ == '__main__':
    # backbone = get_RepVGG_func_by_name('RepVGG-A2')(deploy=True).cuda().eval()
    # head_gender = BaseClassifier(num_classes=2, in_planes=backbone.in_planes, name='gender').cuda().eval()
    # head_age = BaseClassifier(num_classes=4, in_planes=backbone.in_planes, name='age').cuda().eval()
    # head_upcolor = BaseClassifier(num_classes=10, in_planes=backbone.in_planes, name='upcolor').cuda().eval()
    # head_downcolor = BaseClassifier(num_classes=10, in_planes=backbone.in_planes, name='downcolor').cuda().eval()
    #
    # heads = [head_gender, head_age, head_upcolor, head_downcolor]
    # model = FeatClassifier(backbone, heads)

    from utils.get_config import get_config
    from data.unify_attributes import ATTRIBUTE_ALL

    config_file = '/media/aimall/55685ac5-0dd2-4fa5-9b1b-d8fb04f1ccd8/lws/projects/PAR_train/configs/repvgg_a2_market1501.json'
    config = get_config(config_file)
    model = build_network(config['model']['backbone'], ATTRIBUTE_ALL, phase='test')
    model = model.cuda().eval()

    print(model)

    x = torch.randn(1, 3, 288, 288, requires_grad=False).cuda()

    with torch.no_grad():
        outputs = model(x)

    print(outputs)
    # for out in outputs:
    #     print(out)

    # print(model)

    # torch.save(model.state_dict(), 'model_saved/ckpt/model_eval.pth')

        其中，base为一个基础网络结构的类定义：

import torch.nn as nn


class FeatClassifier(nn.Module):

    def __init__(self, backbone, heads):
        super(FeatClassifier, self).__init__()

        assert isinstance(heads, list), "heads must be list"
        self.backbone = backbone
        self.heads = heads
        for head in self.heads:
            self.__setattr__(head.name, head)

    def fresh_params(self):
        params = {}
        for head in self.heads:
            params[head.name] = head.fresh_params()
        return params

    def fresh_params_by_dataset(self, head_name, fresh_params_all):
        return fresh_params_all[head_name]

    def finetune_params(self):
        return self.backbone.parameters()

    def forward(self, x):
        feat_map = self.backbone(x)
        outputs = {}
        for head in self.heads:
            outputs[head.name] = head(feat_map)
        return outputs

        heads的定义为一个小的分类器：

import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules.batchnorm import _BatchNorm


class BaseClassifier(nn.Module):
    def __init__(self, num_classes, in_planes, name):
        super().__init__()
        self.logits = nn.Sequential(
            nn.Linear(in_planes, num_classes),
            # nn.BatchNorm1d(num_classes)
        )
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.name = name

    def fresh_params(self):
        return self.parameters()

    def forward(self, feature):
        # feat = self.avg_pool(feature).view(feature.size(0), -1)
        # feat = self.avg_pool(feature).view(-1, int(feature.size(1)))
        feat = torch.flatten(self.avg_pool(feature), start_dim=1)
        x = self.logits(feat)
        return x

    def warpper_forward(self, feature):
        # feature = torch.flatten(self.avg_pool(feature), start_dim=1)
        x = self.logits(feature)
        x = F.softmax(x, dim=1)
        return x


def initialize_weights(module):
    for m in module.children():
        if isinstance(m, nn.Conv2d):
            n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
            m.weight.data.normal_(0, math.sqrt(2. / n))
        elif isinstance(m, _BatchNorm):
            m.weight.data.fill_(1)
            if m.bias is not None:
                m.bias.data.zero_()
        elif isinstance(m, nn.Linear):
            stdv = 1. / math.sqrt(m.weight.size(1))
            m.weight.data.uniform_(-stdv, stdv)