行人属性识别二:添加新网络训练和自定义数据集训练

序言

上一篇记录了训练过程,但是项目中提供的模型网络都是偏大的,如果想要在边缘设备上部署,还是比较吃力的,所以本文记录如何加入新的网络模型进行训练,以repvgg为例,加入mobilenet、shufflenet等网络同配置。

以及抽取pa100k数据集部分属性作为自己数据集进行训练,对于自主标记的数据集,按照相同格式准备即可。

一、添加新网络repvgg

本项目中用到网络结构其实和分类网络没有太大区别,查看网络的构建,同样是卷积+全连接结构,所以构建起来就简单多了,只需要替换backbone即可。
行人属性识别二:添加新网络训练和自定义数据集训练_第1张图片
在models/backbone中添加repvgg.py文件,文件内容如下,提供的backbone只有两个:RepVGG_A0、RepVGG_A0_m:

import torch.nn as nn
import numpy as np
import torch
import torch.nn.functional as F
from models.registry import BACKBONE

__all__ = ['RepVGG_A0','RepVGG_A0_m']

def conv_bn(in_channels, out_channels, kernel_size, stride, padding, groups=1):
    # 3 x 3的卷积层
    result = nn.Sequential()
    result.add_module('conv', nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
                                                  kernel_size=kernel_size, stride=stride, padding=padding, groups=groups, bias=False))
    result.add_module('bn', nn.BatchNorm2d(num_features=out_channels))
    return result

class RepVGGBlock(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size,
                 stride=1, padding=0, dilation=1, groups=1, padding_mode='zeros', deploy=False):
        super(RepVGGBlock, self).__init__()
        self.deploy = deploy           #   部署
        self.groups = groups
        self.in_channels = in_channels

        assert kernel_size == 3
        assert padding == 1

        padding_11 = padding - kernel_size // 2

        self.nonlinearity = nn.ReLU()     #   激活

        if deploy:
            self.rbr_reparam = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride,
                                      padding=padding, dilation=dilation, groups=groups, bias=True, padding_mode=padding_mode)

        else:
            self.rbr_identity = nn.BatchNorm2d(num_features=in_channels) if out_channels == in_channels and stride == 1 else None
            self.rbr_dense = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups)
            self.rbr_1x1 = conv_bn(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, padding=padding_11, groups=groups)
            # print('RepVGG Block, identity = ', self.rbr_identity)


    def forward(self, inputs):
        if hasattr(self, 'rbr_reparam'):
            return self.nonlinearity(self.rbr_reparam(inputs))

        if self.rbr_identity is None:
            id_out = 0
        else:
            id_out = self.rbr_identity(inputs)

        return self.nonlinearity(self.rbr_dense(inputs) + self.rbr_1x1(inputs) + id_out)


#   This func derives the equivalent kernel and bias in a DIFFERENTIABLE way.
#   You can get the equivalent kernel and bias at any time and do whatever you want,
    #   for example, apply some penalties or constraints during training, just like you do to the other models.
#   May be useful for quantization or pruning.
    def get_equivalent_kernel_bias(self):
        kernel3x3, bias3x3 = self._fuse_bn_tensor(self.rbr_dense)         # 融合3*3卷积和BN
        kernel1x1, bias1x1 = self._fuse_bn_tensor(self.rbr_1x1)           # 融合1*1卷积和BN
        kernelid, biasid = self._fuse_bn_tensor(self.rbr_identity)        # 融合值为1的3*3卷积和BN
        return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid

    def _pad_1x1_to_3x3_tensor(self, kernel1x1):
        if kernel1x1 is None:
            return 0
        else:
            return torch.nn.functional.pad(kernel1x1, [1,1,1,1])

    def _fuse_bn_tensor(self, branch):
        if branch is None:
            return 0, 0
        if isinstance(branch, nn.Sequential):
            kernel = branch.conv.weight
            running_mean = branch.bn.running_mean
            running_var = branch.bn.running_var
            gamma = branch.bn.weight
            beta = branch.bn.bias
            eps = branch.bn.eps
        else:
            assert isinstance(branch, nn.BatchNorm2d)
            if not hasattr(self, 'id_tensor'):
                input_dim = self.in_channels // self.groups
                kernel_value = np.zeros((self.in_channels, input_dim, 3, 3), dtype=np.float32)
                for i in range(self.in_channels):
                    kernel_value[i, i % input_dim, 1, 1] = 1
                self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)
            kernel = self.id_tensor
            running_mean = branch.running_mean
            running_var = branch.running_var
            gamma = branch.weight
            beta = branch.bias
            eps = branch.eps
        std = (running_var + eps).sqrt()
        t = (gamma / std).reshape(-1, 1, 1, 1)
        return kernel * t, beta - running_mean * gamma / std

    def repvgg_convert(self):
        kernel, bias = self.get_equivalent_kernel_bias()
        return kernel.detach().cpu().numpy(), bias.detach().cpu().numpy(),

class RepVGG(nn.Module):

    def __init__(self, num_blocks, num_classes=1000, width_multiplier=None, override_groups_map=None, deploy=False):
        super(RepVGG, self).__init__()

        assert len(width_multiplier) == 4

        self.deploy = deploy
        self.override_groups_map = override_groups_map or dict()

        assert 0 not in self.override_groups_map

        self.in_planes = min(64, int(64 * width_multiplier[0]))
        self.out_planes = int(512 * width_multiplier[3])

        self.stage0 = RepVGGBlock(in_channels=3, out_channels=self.in_planes, kernel_size=3, stride=2, padding=1, deploy=self.deploy)
        self.cur_layer_idx = 1
        self.stage1 = self._make_stage(int(64 * width_multiplier[0]), num_blocks[0], stride=2)
        self.stage2 = self._make_stage(int(128 * width_multiplier[1]), num_blocks[1], stride=2)
        self.stage3 = self._make_stage(int(256 * width_multiplier[2]), num_blocks[2], stride=2)
        self.stage4 = self._make_stage(int(512 * width_multiplier[3]), num_blocks[3], stride=2)
        # self.gap = nn.AdaptiveAvgPool2d(output_size=1)
        # self.linear = nn.Linear(int(512 * width_multiplier[3]), num_classes)


    def _make_stage(self, planes, num_blocks, stride):
        strides = [stride] + [1]*(num_blocks-1)
        blocks = []
        for stride in strides:
            cur_groups = self.override_groups_map.get(self.cur_layer_idx, 1)
            blocks.append(RepVGGBlock(in_channels=self.in_planes, out_channels=planes, kernel_size=3,
                                      stride=stride, padding=1, groups=cur_groups, deploy=self.deploy))
            self.in_planes = planes
            self.cur_layer_idx += 1
        return nn.Sequential(*blocks)

    def forward(self, x):
        out = self.stage0(x)
        out = self.stage1(out)
        out = self.stage2(out)
        out = self.stage3(out)
        out = self.stage4(out)
        # out = self.gap(out)

        # out = out.view(out.size(0), -1)
        # out = self.linear(out)
        return out


optional_groupwise_layers = [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26]
g2_map = {l: 2 for l in optional_groupwise_layers}
g4_map = {l: 4 for l in optional_groupwise_layers}

def create_RepVGG_A0(num_classes=1000,deploy=False):
    return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=num_classes,
                  width_multiplier=[0.75, 0.75, 0.75, 2.5], override_groups_map=None, deploy=deploy)

def create_RepVGG_A0_m(num_classes=1000,deploy=False):
    return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=num_classes,
                  width_multiplier=[0.75, 0.75, 0.75, 1], override_groups_map=None, deploy=deploy)

def create_RepVGG_A0_s(num_classes=1000,deploy=False):
    return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=num_classes,
                  width_multiplier=[0.75, 0.75, 0.5, 1], override_groups_map=None, deploy=deploy)

def create_RepVGG_A1(num_classes=1000,deploy=False):
    return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=num_classes,
                  width_multiplier=[1, 1, 1, 2.5], override_groups_map=None, deploy=deploy)

def create_RepVGG_A2(num_classes=1000,deploy=False):
    return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=num_classes,
                  width_multiplier=[1.5, 1.5, 1.5, 2.75], override_groups_map=None, deploy=deploy)

def create_RepVGG_B0(num_classes=1000,deploy=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=num_classes,
                  width_multiplier=[1, 1, 1, 2.5], override_groups_map=None, deploy=deploy)

def create_RepVGG_B1(num_classes=1000,deploy=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=num_classes,
                  width_multiplier=[2, 2, 2, 4], override_groups_map=None, deploy=deploy)

def create_RepVGG_B1g2(deploy=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[2, 2, 2, 4], override_groups_map=g2_map, deploy=deploy)

def create_RepVGG_B1g4(num_classes=1000,deploy=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=num_classes,
                  width_multiplier=[2, 2, 2, 4], override_groups_map=g4_map, deploy=deploy)

def create_RepVGG_B2(num_classes=1000,deploy=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=num_classes,
                  width_multiplier=[2.5, 2.5, 2.5, 5], override_groups_map=None, deploy=deploy)

def create_RepVGG_B2g2(deploy=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[2.5, 2.5, 2.5, 5], override_groups_map=g2_map, deploy=deploy)

def create_RepVGG_B2g4(deploy=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[2.5, 2.5, 2.5, 5], override_groups_map=g4_map, deploy=deploy)


def create_RepVGG_B3(num_classes=1000,deploy=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=num_classes,
                  width_multiplier=[3, 3, 3, 5], override_groups_map=None, deploy=deploy)


def create_RepVGG_B3g2(deploy=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[3, 3, 3, 5], override_groups_map=g2_map, deploy=deploy)

def create_RepVGG_B3g4(deploy=False):
    return RepVGG(num_blocks=[4, 6, 16, 1], num_classes=1000,
                  width_multiplier=[3, 3, 3, 5], override_groups_map=g4_map, deploy=deploy)


func_dict = {
'RepVGG-A0': create_RepVGG_A0,
'RepVGG-A0_s': create_RepVGG_A0_s,
'RepVGG-A0_m': create_RepVGG_A0_m,
'RepVGG-A1': create_RepVGG_A1,
'RepVGG-A2': create_RepVGG_A2,
'RepVGG-B0': create_RepVGG_B0,
'RepVGG-B1': create_RepVGG_B1,
'RepVGG-B1g2': create_RepVGG_B1g2,
'RepVGG-B1g4': create_RepVGG_B1g4,
'RepVGG-B2': create_RepVGG_B2,
'RepVGG-B2g2': create_RepVGG_B2g2,
'RepVGG-B2g4': create_RepVGG_B2g4,
'RepVGG-B3': create_RepVGG_B3,
'RepVGG-B3g2': create_RepVGG_B3g2,
'RepVGG-B3g4': create_RepVGG_B3g4,
}
def get_RepVGG_func_by_name(name):
    return func_dict[name]


#   Use like this:
#   train_model = create_RepVGG_A0(deploy=False)
#   train train_model
#   deploy_model = repvgg_convert(train_model, create_RepVGG_A0, save_path='repvgg_deploy.pth')
def repvgg_model_convert(model:torch.nn.Module, build_func, save_path=None,num_classes=1000):
    converted_weights = {}
    for name, module in model.named_modules():
        if hasattr(module, 'repvgg_convert'):
            kernel, bias = module.repvgg_convert()
            converted_weights[name + '.rbr_reparam.weight'] = kernel
            converted_weights[name + '.rbr_reparam.bias'] = bias
        elif isinstance(module, torch.nn.Linear):
            converted_weights[name + '.weight'] = module.weight.detach().cpu().numpy()
            converted_weights[name + '.bias'] = module.bias.detach().cpu().numpy()
        else:
            print(name, type(module))
    del model

    deploy_model = build_func(num_classes=num_classes,deploy=True)
    for name, param in deploy_model.named_parameters():
        print('deploy param: ', name, param.size(), np.mean(converted_weights[name]))
        param.data = torch.from_numpy(converted_weights[name]).float()

    if save_path is not None:
        torch.save(deploy_model.state_dict(), save_path,_use_new_zipfile_serialization=False)

    return deploy_model

@BACKBONE.register("repvgg_a0")
def RepVGG_A0(model_path = "/home/cai/project/Rethinking_of_PAR/model/RepVGG-A0-train.pth"):
    model = create_RepVGG_A0()
    if model_path is not None:
        pretrained_params = torch.load(model_path,map_location=torch.device("cpu"))
        # pretrained_params= \
        #     {k: v for k, v in pretrained_params.items() if
        #      k in net.state_dict().keys() and net.state_dict()[k].numel() == v.numel()}
        model.load_state_dict(pretrained_params, strict=False)

    return model

@BACKBONE.register("repvgg_a0_m")
def RepVGG_A0_m(model_path = "/home/cai/project/Rethinking_of_PAR/model/RepVGG-A0-train.pth"):
    model = create_RepVGG_A0_m()
    if model_path is not None:
        pretrained_params = torch.load(model_path,map_location=torch.device("cpu"))
        # pretrained_params= \
        #     {k: v for k, v in pretrained_params.items() if
        #      k in net.state_dict().keys() and net.state_dict()[k].numel() == v.numel()}
        model.load_state_dict(pretrained_params, strict=False)

    return model

需要注意的是这两个部分,在初始化网络前要加载好预训练模型,其他网络也是类似:
行人属性识别二:添加新网络训练和自定义数据集训练_第2张图片
然后在models/model_factory.py中添加,这两个网络最后层输出的对应通道数:
行人属性识别二:添加新网络训练和自定义数据集训练_第3张图片
在train.py中导入
行人属性识别二:添加新网络训练和自定义数据集训练_第4张图片
最后修改配置文件,添加repvgg_a0网络,就可以运行了:
行人属性识别二:添加新网络训练和自定义数据集训练_第5张图片

二、自定义数据集

pa100k数据集中包含了26种属性,有一些属性对我来说没有用,不想训练这么多,所以抽取其中18中属性进行训练,那么该如何准备呢?

首先,先通过以下脚本将pa100k的mat标签文件转换为txt形式:

import pandas as pd
import scipy
from scipy import io

def mat2txt(data, key):
    subdata = data[key]
    dfdata = pd.DataFrame(subdata)
    dfdata.to_csv("/home/cai/data/PA100K/%s.txt" % key, index=False)

if __name__ == "__main__":
    data = scipy.io.loadmat("/home/cai/data/PA100K/annotation.mat")
    key_list = ["attributes", "test_images_name", "test_label",
                "train_images_name", "train_label",
                "val_images_name", "val_label"]
    for key in key_list:
        mat2txt(data, key)

得到如下几个文件:
行人属性识别二:添加新网络训练和自定义数据集训练_第6张图片
在通过如下文件对属性进行剔除,并声称新的txt文件:

# 根据txt生成标签
import os

txts = [["train_images_name", "train_label"],["test_images_name", "test_label"],["val_images_name", "val_label"]]
txt_path = "/home/cai/data/PA100K"

for txt_list in txts:
    file1 = open(os.path.join(txt_path,txt_list[0]+".txt"),"r")
    file2 = open(os.path.join(txt_path,txt_list[1]+".txt"),"r")
    save_file = open(txt_list[0].split("_")[0]+".txt","w")
    label1_list = []
    label2_list = []
    for line1 in file1.readlines():
        label1_list.append(line1)
    for line2 in file2.readlines():
        label2_list.append(line2)
    for i,label in enumerate(label1_list):
        if i ==0 :
            continue
        label1 = label1_list[i].split('\n')[0].split('\'')
        label2 = label2_list[i].split(',')
        for i,ind in enumerate([12,15,16,17,18,19,20,25]):                   # 这是要剔除的相应属性索引
            label2.pop(ind-i)
        label2 = ",".join(label2)
        save_file.write("/home/cai/data/PA100K/PA100k/data/"+label1[1]+"\t"+label2+"\n")
    file1.close()
    file2.close()
    save_file.close()

得到三个文件train.txt、test.txt、val.txt,内容如下,如果是自己标注的数据集,想办法转成这个格式即可:
行人属性识别二:添加新网络训练和自定义数据集训练_第7张图片

然后根据dataset/pedes_attr/preprocess/format_pa100k.py编写format_mydata.py文件,在生成pkl之前需要对数据集属性进行重排,重排规则看dataset/pedes_attr/annotation.md,其实不排也可以,不排的话将generate_data_description函数中label_txt,reorder=Flase即可,相应的修改位置已经在代码中标明:

import os
import numpy as np
import random
import pickle

from easydict import EasyDict
from scipy.io import loadmat

np.random.seed(0)
random.seed(0)

classes_name = ['Female','AgeOver60','Age18-60','AgeLess18','Front','Side','Back','Hat','Glasses','HandBag','ShoulderBag','Backpack'
                ,'ShortSleeve','LongSleeve','LongCoat','Trousers','Shorts','Skirt&Dress']                  # 属性类别

group_order = [7, 8, 12, 13, 14, 15, 16, 17, 9, 10, 11, 1, 2, 3, 0, 4, 5, 6]                         #  需要对新属性进行重排,属性重排后的索引顺序(不懂的看上一篇)
# clas_name = ['Hat','Glasses','ShortSleeve','LongSleeve','LongCoat','Trousers','Shorts','Skirt&Dress','HandBag','ShoulderBag','Backpack'
#              ,'AgeOver60','Age18-60','AgeLess18','Female','Front','Side','Back']

def make_dir(path):
    if os.path.exists(path):
        pass
    else:
        os.mkdir(path)

def generate_data_description(save_dir, label_txt,reorder):
    """
    create a dataset description file, which consists of images, labels
    """
    image_name = []
    image_label = []

    file = open(label_txt,"r")
    for line in file.readlines():
        name = line.split("\t")[0]
        label = line.split('\t')[1].split('\n')[0].split(',')
        label = list(map(int, label))
        image_name.append(name)
        image_label.append(label)

    dataset = EasyDict()
    dataset.description = 'pa100k'
    dataset.reorder = 'group_order'
    dataset.root = os.path.join(save_dir, 'data')

    dataset.image_name = image_name
    dataset.label = np.array(image_label)
    dataset.attr_name = classes_name

    dataset.label_idx = EasyDict()
    dataset.label_idx.eval = list(range(len(classes_name)))

    if reorder:
        dataset.label_idx.eval = group_order

    dataset.partition = EasyDict()
    dataset.partition.train = np.arange(0, 80000)      # np.array(range(80000))                             # 数据集数量划分,自己根据自己数据集来,这是pa100k的划分
    dataset.partition.val = np.arange(80000, 90000)    # np.array(range(80000, 90000))
    dataset.partition.test = np.arange(90000, 100000)  # np.array(range(90000, 100000))
    dataset.partition.trainval = np.arange(0, 90000)   # np.array(range(90000))

    dataset.weight_train = np.mean(dataset.label[dataset.partition.train], axis=0).astype(np.float32)
    dataset.weight_trainval = np.mean(dataset.label[dataset.partition.trainval], axis=0).astype(np.float32)

    with open(os.path.join(save_dir, 'dataset_all.pkl'), 'wb+') as f:
        pickle.dump(dataset, f)

if __name__ == "__main__":
    # save_dir = '/mnt/data1/jiajian/datasets/attribute/PA100k/'
    save_dir = '/home/cai/data/PA100K/MyData/'                                             # 数据集图片存放路径   MyData/data   MyData/dataset_all.pkl
    label_txt = "/home/cai/project/Rethinking_of_PAR/data/MyData/label.txt"                # train.txt test.txt val.txt 合并成label.txt
    generate_data_description(save_dir, label_txt,reorder=True)

修改完成后需要到该文件下添加自己数据集名字,我的数据集取名为MyData:
行人属性识别二:添加新网络训练和自定义数据集训练_第8张图片
在configs中新建mydata.yaml,内容根据pa100k的来修改,示例如下,加入了新的网络和数据集:
行人属性识别二:添加新网络训练和自定义数据集训练_第9张图片
最后运行代码训练:

python train.py --cfg configs/pedes_baseline/mydata.yaml

行人属性识别二:添加新网络训练和自定义数据集训练_第10张图片

你可能感兴趣的:(行人属性识别,网络,深度学习,人工智能)