当前位置：首页 > news >正文

基于Pytorch框架的深度学习HRnet网络人像语义分割系统源码

news 2025/5/11 3:02:06

第一步：准备数据

头发分割数据，总共有5711张图片，里面的像素值为0和1，所以看起来全部是黑的，不影响使用

第二步：搭建模型

计算机视觉领域有很多任务是位置敏感的，比如目标检测、语义分割、实例分割等等。为了这些任务位置信息更加精准，很容易想到的做法就是维持高分辨率的feature map，事实上HRNet之前几乎所有的网络都是这么做的，通过下采样得到强语义信息，然后再上采样恢复高分辨率恢复位置信息(如下图所示)，然而这种做法，会导致大量的有效信息在不断的上下采样过程中丢失。而HRNet通过并行多个分辨率的分支，加上不断进行不同分支之间的信息交互，同时达到强语义信息和精准位置信息的目的。

recover high resolution

思路在当时来讲，不同分支的信息交互属于很老套的思路（如FPN等），我觉得最大的创新点还是能够从头到尾保持高分辨率，而不同分支的信息交互是为了补充通道数减少带来的信息损耗，这种网络架构设计对于位置敏感的任务会有奇效。

第三步：代码

1）损失函数为：交叉熵损失函数

2）网络代码：

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as Ffrom .backbone import BN_MOMENTUM, hrnet_classificationclass HRnet_Backbone(nn.Module):def __init__(self, backbone = 'hrnetv2_w18', pretrained = False):super(HRnet_Backbone, self).__init__()self.model    = hrnet_classification(backbone = backbone, pretrained = pretrained)del self.model.incre_modulesdel self.model.downsamp_modulesdel self.model.final_layerdel self.model.classifierdef forward(self, x):x = self.model.conv1(x)x = self.model.bn1(x)x = self.model.relu(x)x = self.model.conv2(x)x = self.model.bn2(x)x = self.model.relu(x)x = self.model.layer1(x)x_list = []for i in range(2):if self.model.transition1[i] is not None:x_list.append(self.model.transition1[i](x))else:x_list.append(x)y_list = self.model.stage2(x_list)x_list = []for i in range(3):if self.model.transition2[i] is not None:if i < 2:x_list.append(self.model.transition2[i](y_list[i]))else:x_list.append(self.model.transition2[i](y_list[-1]))else:x_list.append(y_list[i])y_list = self.model.stage3(x_list)x_list = []for i in range(4):if self.model.transition3[i] is not None:if i < 3:x_list.append(self.model.transition3[i](y_list[i]))else:x_list.append(self.model.transition3[i](y_list[-1]))else:x_list.append(y_list[i])y_list = self.model.stage4(x_list)return y_listclass HRnet(nn.Module):def __init__(self, num_classes = 21, backbone = 'hrnetv2_w18', pretrained = False):super(HRnet, self).__init__()self.backbone       = HRnet_Backbone(backbone = backbone, pretrained = pretrained)last_inp_channels   = np.int(np.sum(self.backbone.model.pre_stage_channels))self.last_layer = nn.Sequential(nn.Conv2d(in_channels=last_inp_channels, out_channels=last_inp_channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(last_inp_channels, momentum=BN_MOMENTUM),nn.ReLU(inplace=True),nn.Conv2d(in_channels=last_inp_channels, out_channels=num_classes, kernel_size=1, stride=1, padding=0))def forward(self, inputs):H, W = inputs.size(2), inputs.size(3)x = self.backbone(inputs)# Upsamplingx0_h, x0_w = x[0].size(2), x[0].size(3)x1 = F.interpolate(x[1], size=(x0_h, x0_w), mode='bilinear', align_corners=True)x2 = F.interpolate(x[2], size=(x0_h, x0_w), mode='bilinear', align_corners=True)x3 = F.interpolate(x[3], size=(x0_h, x0_w), mode='bilinear', align_corners=True)x = torch.cat([x[0], x1, x2, x3], 1)x = self.last_layer(x)x = F.interpolate(x, size=(H, W), mode='bilinear', align_corners=True)return x

第四步：统计一些指标（训练过程中的loss和miou）