轻松学Pytorch – 年龄与性别预测-技术圈

点击上方“小白学视觉”，选择加"星标"或“置顶”

重磅干货，第一时间送达

大家好，上周太忙，没有更新Pytorch轻松学系列文章，但是我还是会坚定的继续走下去的，所谓有始有终，这个系列我会一直坚持写下去，希望大家继续支持我，积极给我反馈，当然也感谢大家的信任与点赞支持。

本文主要是基于公开数据集，完成了一个人脸的年龄与性别预测网络模型，以及模型训练与导出使用、本篇主要讲述的知识点有以下：

如何实现卷积神经网络的多任务不同输出
如何同时实现分类跟回归预测
基于人脸年龄与性别的公开数据的数据制作
使用多任务网络实现推理预测

数据集

本文使用的数据集来自这里

https://susanqq.github.io/UTKFace/

我使用的是已经对齐跟剪切之后的人脸数据，超过2W多张的标注数据，标注信息如下：

[age]_[gender]_[race]_[date&time].jpg

文件名称格式就是每张图像的标注信息

Age表示年龄，范围在0~116岁之间
Gender表示性别，0表示男性，1表示女性
Race表示人种，

基于Pytorch的dataset超类，重新完成了我的自定义数据集，代码如下：

class AgeGenderDataset(Dataset):     def __init__(self, root_dir):         self.transform = transforms.Compose([transforms.ToTensor()])         img_files = os.listdir(root_dir)         nums_ = len(img_files)         # age: 0 ~116, 0 :male, 1 :female         self.ages = []         self.genders = []         self.images = []         index = 0         for file_name in img_files:             age_gender_group = file_name.split("_")             age_ = age_gender_group[0]             gender_ = age_gender_group[1]             self.genders.append(np.float32(gender_))             self.ages.append(np.float32(age_)/max_age)             self.images.append(os.path.join(root_dir, file_name))             index += 1
     def __len__(self):         return len(self.images)
     def num_of_samples(self):         return len(self.images)
     def __getitem__(self, idx):         if torch.is_tensor(idx):             idx = idx.tolist()             image_path = self.images[idx]         else:             image_path = self.images[idx]         img = cv.imread(image_path)  # BGR order         h, w, c = img.shape         # rescale         img = cv.resize(img, (64, 64))         img = (np.float32(img) /255.0 - 0.5) / 0.5         # H, W C to C, H, W         img = img.transpose((2, 0, 1))         sample = {'image': torch.from_numpy(img), 'age': self.ages[idx], 'gender': self.genders[idx]}         return sample

网络模型结构

卷积层作为基础模型，在卷积层之后使用最大全局池化，完成降维操作，然后再此基础上分为两路的全链接层，分布预测年龄跟性别分类，模型结构详细信息如下：

MyMulitpleTaskNet(  (cnn_layers): Sequential(    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))    (1): ReLU()    (2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)    (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)    (4): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))    (5): ReLU()    (6): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)    (7): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)    (8): Conv2d(64, 96, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))    (9): ReLU()    (10): BatchNorm2d(96, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)    (11): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)    (12): Conv2d(96, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))    (13): ReLU()    (14): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)    (15): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)    (16): Conv2d(128, 196, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))    (17): ReLU()    (18): BatchNorm2d(196, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)    (19): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)  )  (global_max_pooling): AdaptiveMaxPool2d(output_size=(1, 1))  (age_fc_layers): Sequential(    (0): Linear(in_features=196, out_features=25, bias=True)    (1): ReLU()    (2): Linear(in_features=25, out_features=1, bias=True)    (3): Sigmoid()  )  (gender_fc_layers): Sequential(    (0): Linear(in_features=196, out_features=25, bias=True)    (1): ReLU()    (2): Linear(in_features=25, out_features=2, bias=True)  ))

其中年龄是一个数值结果，所以通过回归来预测，这里使用sigmoid，需要注意的是在制作数据集的时候需要把年龄归一化到0~1之间。性别是二分类预测，使用softmax，实现预测。损失函数选择，对于年龄预测使用MSE，对于分类预测使用了交叉熵损失，总的损失函数是二值之和、权重相同

模型训练

模型训练的输入图像格式为BGR、NCHW = Nx3x64x64、每个批次N=16个样本，代码实现如下：

if train_on_gpu:     model.cuda()
 ds = AgeGenderDataset("D:/python/pytorch_tutorial/UTKFace/") num_train_samples = ds.num_of_samples() bs = 16 dataloader = DataLoader(ds, batch_size=bs, shuffle=True)
 # 训练模型的次数 num_epochs = 25 # optimizer = torch.optim.SGD(model.parameters(), lr=0.001) optimizer = torch.optim.Adam(model.parameters(), lr=1e-2) model.train()
 # 损失函数 mse_loss = torch.nn.MSELoss() cross_loss = torch.nn.CrossEntropyLoss() index = 0 for epoch in  range(num_epochs):     train_loss = 0.0     for i_batch, sample_batched in enumerate(dataloader):         images_batch, age_batch, gender_batch = \             sample_batched['image'], sample_batched['age'], sample_batched['gender']         if train_on_gpu:             images_batch, age_batch, gender_batch = images_batch.cuda(), age_batch.cuda(), gender_batch.cuda()         optimizer.zero_grad()
         # forward pass: compute predicted outputs by passing inputs to the model         m_age_out_, m_gender_out_ = model(images_batch)         age_batch = age_batch.view(-1, 1)         gender_batch = gender_batch.long()
         # calculate the batch loss         loss = mse_loss(m_age_out_, age_batch) + cross_loss(m_gender_out_, gender_batch)
         # backward pass: compute gradient of the loss with respect to model parameters         loss.backward()
         # perform a single optimization step (parameter update)         optimizer.step()
         # update training loss         train_loss += loss.item()         if index % 100 == 0:             print('step: {} \tTraining Loss: {:.6f} '.format(index, loss.item()))         index += 1
         # 计算平均损失     train_loss = train_loss / num_train_samples
     # 显示训练集与验证集的损失函数     print('Epoch: {} \tTraining Loss: {:.6f} '.format(epoch, train_loss))
 # save model model.eval() torch.save(model, 'age_gender_model.pt')

网络模型结构

基于训练模型，实现人脸年龄与性别预测，人脸检测模型使用OpenCV DNN模块，代码实现如下：

 1def video_landmark_demo():
 2    cnn_model = torch.load("./age_gender_model.pt")
 3    print(cnn_model)
 4    # capture = cv.VideoCapture(0)
 5    capture = cv.VideoCapture("D:/images/video/example_dsh.mp4")
 6
 7    # load tensorflow model
 8    net = cv.dnn.readNetFromTensorflow(model_bin, config=config_text)
 9    while True:
10        ret, frame = capture.read()
11        if ret is not True:
12            break
13        frame = cv.flip(frame, 1)
14        h, w, c = frame.shape
15        blobImage = cv.dnn.blobFromImage(frame, 1.0, (300, 300), (104.0, 177.0, 123.0), False, False);
16        net.setInput(blobImage)
17        cvOut = net.forward()
18        # 绘制检测矩形
19        for detection in cvOut[0,0,:,:]:
20            score = float(detection[2])
21            if score > 0.5:
22                left = detection[3]*w
23                top = detection[4]*h
24                right = detection[5]*w
25                bottom = detection[6]*h
26
27                # roi and detect landmark
28                roi = frame[np.int32(top):np.int32(bottom),np.int32(left):np.int32(right),:]
29                rw = right - left
30                rh = bottom - top
31                img = cv.resize(roi, (64, 64))
32                img = (np.float32(img) / 255.0 - 0.5) / 0.5
33                img = img.transpose((2, 0, 1))
34                x_input = torch.from_numpy(img).view(1, 3, 64, 64)
35                age_, gender_ = cnn_model(x_input.cuda())
36                predict_gender = torch.max(gender_, 1)[1].cpu().detach().numpy()[0]
37                gender = "Male"
38                if predict_gender == 1:
39                    gender = "Female"
40                predict_age = age_.cpu().detach().numpy()*116.0
41                print(predict_gender, predict_age)
42
43                # 绘制
44                cv.putText(frame, ("gender: %s, age:%d"%(gender, int(predict_age[0][0]))), (int(left), int(top)-15), cv.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 1)
45                cv.rectangle(frame, (int(left), int(top)), (int(right), int(bottom)), (255, 0, 0), thickness=2)
46                # cv.putText(frame, "score:%.2f"%score, (int(left), int(top)), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
47                c = cv.waitKey(10)
48                if c == 27:
49                    break
50                cv.imshow("face detection + landmark", frame)
51
52    cv.waitKey(0)
53    cv.destroyAllWindows()
54
55
56if __name__ == "__main__":
57    video_landmark_demo()

运行结果如下：

下载1：OpenCV-Contrib扩展模块中文版教程

在「小白学视觉」公众号后台回复：扩展模块中文教程，即可下载全网第一份OpenCV扩展模块教程中文版，涵盖扩展模块安装、SFM算法、立体视觉、目标跟踪、生物视觉、超分辨率处理等二十多章内容。

下载2：Python视觉实战项目52讲

在「小白学视觉」公众号后台回复：Python视觉实战项目，即可下载包括图像分割、口罩检测、车道线检测、车辆计数、添加眼线、车牌识别、字符识别、情绪检测、文本内容提取、面部识别等31个视觉实战项目，助力快速学校计算机视觉。

下载3：OpenCV实战项目20讲

在「小白学视觉」公众号后台回复：OpenCV实战项目20讲，即可下载含有20个基于OpenCV实现20个实战项目，实现OpenCV学习进阶。

交流群

欢迎加入公众号读者群一起和同行交流，目前有SLAM、三维视觉、传感器、自动驾驶、计算摄影、检测、分割、识别、医学影像、GAN、算法竞赛等微信群（以后会逐渐细分），请扫描下面微信号加群，备注：”昵称+学校/公司+研究方向“，例如：”张三 + 上海交大 + 视觉SLAM“。请按照格式备注，否则不予通过。添加成功后会根据研究方向邀请进入相关微信群。请勿在群内发送广告，否则会请出群，谢谢理解~