PSPNet_PSPNET自己的数据集「建议收藏」

PSPNnet结构介绍

流程:

输入图像进行letterboximage,将它改变为473×473的一个大小-传入主干特征提取网络进行特征提取(Resnet mobilNet)-获得特征层，将Feature Map分为两类，一类直接传入后方，另一类传入psp模块中，使用不同大小的池化层对输入进来的池化层进行平均池化，进行上采样后与直接传入后方的Feature Map进行堆叠-将输入进来的特征层划分成6×6，3×3，2×2，1×1的网格-堆叠后会对得到的图的特征进行卷积和上采样，会对每一个像素点进行分类(输出结果shape与输入相同473×473)

PSP结构功能：

将获取到的特征层划分成不同大小的网格，每个网格内部各自进行平均池化。实现聚合不同区域的上下文信息，从而提高获取全局信息的能力。

红色：将输入进来的特征层整个进行平均池化。

橙色：将输入进来的特征层划分为2×2个子区域，然后对每个子区域进行平均池化。

蓝色：将输入进来的特征层划分为3×3个子区域，然后对每个子区域进行平均池化。

绿色：将输入进来的特征层划分为6×6个子区域，然后对每个子区域进行平均池化。

主干提取网络(mobilenetv2)

Inverted residuals

主干(左边)部分:

1. 对输入的特征层利用1×1的卷积进行升维
2. 升维过后进行标准化和激活函数
3. 利用3×3的可分离卷积进行特征提取
4. 再进行标准化和激活函数
5. 利用1×1的卷积进行降维(调整主干道的通道数，使通道数与残差部分相同)
6. 再进行标准化

代码(主干部分)

class InvertedResidual(nn.Module):
    def __init__(self,inp,oup,stride,expand_ratio):
        super(InvertedResidual,self).__init__()
        self.stride = stride
        assert stride in [1,2]
        
        hidden_dim = round(inp*expand_ratio)
        self.use_res_connect = self.stride == 1 and inp == oup
	if expand_ratio == 1:#判断图片需不需要进行1x1的升维，若等于1则不需要进行升维直接进行3x3的特征提取
    	self.conv == nn.sequential(
    		#dw深度可分离卷积
        	nn.Conv2d(hidden_dim,hidden_dim,3,stride,1,groups=hidden_dim,bias=False),
        	BatchNorm2d(hidden_dim),
        	nn.ReLU6(inplace=True),
        	#pw-linear1x1降维
        	nn.Conv2d(hidden_dim,oup,1,1,0,bias=False),
        	BatchNorm2d(oup),    
    	)
	else:
    	self.conv = nn.Sequential(
    		#pw1x1卷积升维
        	nn.Conv2d(inp,hidden_dim,1,1,0,bias=False),
        	BatchNorm2d(hidden_dim),
        	nn.ReLU6(inplace=True),
        	#dw
        	nn.Conv2d(hidden_dim,hidden_dim,3,stride,1,groups=hidden_dim,bias=False),
        	BatchNorm2d(hidden_dim),
        	nn.ReLU6(inplace=True),
        	#pw-linear
        	nn.Conv2d(hidden_dim,oup,1,1,0,bias=False),
        	BatchNorm2d(oup),#标准化
    	)
	def forward(self,x):
    	if self.use_res_connect:#若使用残差部分为True
        	return x + self.conv(x)#则会将主干部分输出结果与残差部分进行相加
        else:
            	return self.conv(x)#若不使用则将主干部分结果直接输出

利用Inverted residuals完成mobilenetv2的构建

class MobileNetV2(nn.Module):
    def __init__(self,n_class=1000,input_size=224,width_mult=1.):
        super(MobileNetV2,self).__init__()
        block = InvertedResidual
        input_channel = 32
        last_channel = 1280
        
        interverted_residual_setting=[
            #t,c,n,s 列表，每行都可看作大机构快，每一个都由Inverted residuals组成
            #473,473,3 ->237,237,32 第一次卷积标准化和激活函数，步长为2
            #237,237,32 -> 237，237,16 第一个大结构快通道数是16，故不会改变高和宽，只改变通道数
            [1,16,1,1],#只有这个第一个参数1，不会进行1x1卷积升维
            #步长是2 通道数是24
            #237,237,16, ->119,119,24
            [6,24,2,2],#第二个参数oup最终输出通道数
            # 119,119,24 -> 60,60,32
            [6,32,3,2],#第三个参数是Inverted residuals的数量
            # 60,60,32 -> 30,30,64
            [6,64,4,2],#第四个参数，步长，判断模块是否对输入的特征层进行压缩
            # 30,30,64 -> 30,30,96
            [6,96,3,1],
            # 30,30,96 -> 15,15,160
            [6,160,3,2],
            # 15,15,160 -> 15,15,320
            [6,320,1,1],
        ]
        assert input_size % 32 == 0
        #建立stem层
        input_channel = int(input_channel * width_mult)
        self.last_channel = int(last_channel*width_mult) if width_mult > 1.0 else last_channel
        self.features = [conv_bn(3,input_channel,2)]#创建features变量，代表所有层集合
        #conv_bn卷积标准化+激活函数
        
        #根据上述列表进行循环，构建mobilenetv2的结构
        for t,c,n,s in interverted_residual_setting:#分别取4个参数进行循环
            output_channel = int(c * width_mult)
            for i in range(n):#对n进行循环，代表大结构块中Inverted residuals的重复次数
                if i == 0:#若循环次数是第一次，则考虑是否对输入进来的特征层进行高和宽的压缩
                    self.feature.append(block(input_channel,output_channel,s,expand_ratio=t))
                else:
                    self.feature.append(block(input_channel,output_channel,1,expand_ratio=t))
                    input_channel = output_channel
                    

加强特征提取网络-PSP模块的构建

PSPNet所使用的加强特征提取结构是PSP模块。

PSP结构的做法是将**获取到的特征层划分成不同大小的区域，每个区域内部各自进行平均池化。**实现聚合不同区域的上下文信息，从而提高获取全局信息的能力。

在PSPNet中，PSP结构典型情况下，会将输入进来的特征层划分成6×6，3×3，2×2，1×1的区域，然后每个区域内部各自进行平均池化。

假设PSP结构输入进来的特征层为30x30x320，此时这个特征层的高和宽均为30，如果我们要将这个特征层划分成6×6的区域，只需要使得平均池化的步长stride=30/6=5和kernel_size=30/6=5就行了，此时的平均池化相当于将特征层划分成6×6的区域，每个区域内部各自进行平均池化。

当PSP结构输入进来的特征层为30x30x320时，PSP结构的具体构成如下。

代码

class _PSPModule(nn.Module):
    def __init__(self, in_channels, pool_sizes, norm_layer):#pool_size大小为1,2,3,6分别对应PSP经典结构的各区域
        super(_PSPModule, self).__init__()
        out_channels = in_channels // len(pool_sizes)
        self.stages = nn.ModuleList([self._make_stages(in_channels, out_channels, pool_size, norm_layer) 
                                                        for pool_size in pool_sizes])#对pool_size进行循环
        #self.bottleneck将堆叠后的特征层进行调整 in_channels就是Feature Map的通道数，out_channels是池化后的通道数，pool_sizes的长度四个平均池化相加后的结果
        self.bottleneck = nn.Sequential(
            nn.Conv2d(in_channels+(out_channels * len(pool_sizes)), out_channels, 
                                    kernel_size=3, padding=1, bias=False),
            norm_layer(out_channels),
            nn.ReLU(inplace=True),
            nn.Dropout2d(0.1)
        )

    def _make_stages(self, in_channels, out_channels, bin_sz, norm_layer):#定义平均池化
        prior = nn.AdaptiveAvgPool2d(output_size=bin_sz)#自适应平均池化层，可直接指定输入特征层输出的shape为多少
        conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)#1x1卷积进行通道手调整
        bn = norm_layer(out_channels)#卷积标准化
        relu = nn.ReLU(inplace=True)#激活函数
        return nn.Sequential(prior, conv, bn, relu)#将这几层保存在self.stages里
    
    def forward(self, features):
        h, w = features.size()[2], features.size()[3]
        pyramids = [features]#将特征层输入到列表中
        #stage将特征层划分成6x6，3x3,2x2,1x1的区域再进行平均池化
        pyramids.extend([F.interpolate(stage(features), size=(h, w), mode='bilinear', 
                                        align_corners=True) for stage in self.stages])
        #F.interpolate将获取到的特征Resize成Feature的大小，之后才可进行堆叠
        output = self.bottleneck(torch.cat(pyramids, dim=1))
        #堆叠
        return output

PSPhead-利用特征获得最终的预测结果

利用加强特征获得预测结果

1、利用一个3×3卷积对特征进行整合。

2、利用一个1×1卷积进行通道调整，调整成Num_Classes。

3、利用resize进行上采样使得最终输出层，宽高和输入图片一样。

class PSPNet(nn.Module):
    def __init__(self, num_classes, downsample_factor, backbone="resnet50", pretrained=True, aux_branch=True):
        super(PSPNet, self).__init__()
        norm_layer = nn.BatchNorm2d
        if backbone=="resnet50":
            self.backbone = Resnet(downsample_factor, pretrained)
            aux_channel = 1024
            out_channel = 2048
        elif backbone=="mobilenet":
            self.backbone = MobileNetV2(downsample_factor, pretrained)
            aux_channel = 96
            out_channel = 320
        else:
            raise ValueError('Unsupported backbone - `{}`, Use mobilenet, resnet50.'.format(backbone))

        self.master_branch = nn.Sequential(
            _PSPModule(out_channel, pool_sizes=[1, 2, 3, 6], norm_layer=norm_layer),
            nn.Conv2d(out_channel//4, num_classes, kernel_size=1)
        )#1x1的卷积通道数是num_classes相当于对每一个特征点进行分类

        self.aux_branch = aux_branch

        if self.aux_branch:#构建辅助训练分支
            self.auxiliary_branch = nn.Sequential(#首先获得辅助特征层，一般选取导数第三个大结构快
                nn.Conv2d(aux_channel, out_channel//8, kernel_size=3, padding=1, bias=False),#利用3x3卷积进行特征整合，再进行通道数调整
                norm_layer(out_channel//8),
                nn.ReLU(inplace=True),
                nn.Dropout2d(0.1),
                nn.Conv2d(out_channel//8, num_classes, kernel_size=1)
            )#再利用1x1的卷积进行通道数的调整，将通道数调整成num_classes，相当于将特征点进行分类，再进行resize到输入图片大小

        self.initialize_weights(self.master_branch)

    def forward(self, x):
        input_size = (x.size()[2], x.size()[3])
        x_aux, x = self.backbone(x)

        output = self.master_branch(x)
        output = F.interpolate(output, size=input_size, mode='bilinear', align_corners=True)#对分类后的特征点进行Rsize，Rsize到输入图层的大小，再辅助训练分支输出和标签进行比较
        if self.aux_branch:
            output_aux = self.auxiliary_branch(x_aux)
            output_aux = F.interpolate(output_aux, size=input_size, mode='bilinear', align_corners=True)
            return output_aux, output
        else:
            return output

    def initialize_weights(self, *models):
        for model in models:
            for m in model.modules():
                if isinstance(m, nn.Conv2d):
                    nn.init.kaiming_normal_(m.weight.data, nonlinearity='relu')
                elif isinstance(m, nn.BatchNorm2d):
                    m.weight.data.fill_(1.)
                    m.bias.data.fill_(1e-4)
                elif isinstance(m, nn.Linear):
                    m.weight.data.normal_(0.0, 0.0001)
                    m.bias.data.zero_()

预测过程

from pspnet import PSPNet
from PIL import Image

pspnet = PSPNet()#主要方法detect_image
#def detect_image(self,image):
	#old_img = copy.deepcopy(image)对图片进行备份，完成图片分割后可对图片进行绘制
	#orininal_h = np.array(image).shape[0] 计算原图的高
    #orininal_w = np.array(image).shape[1] 计算原图的宽
    
    #image,nw,nh = self.letterbox_image(image,(self.model_image_size[1],self.model_image_size[0])) 进行不失真的resize
    #images = [np.array(image)/255] 对图像归一化
    #images = np.transpose(images,[0,3,1,2]) 将通道调整为第二维度

    #with torch.no_grad():
    	#images = variable(torch.from_numpy(images).type(torch.FloatTensor))将图片转化成Tensor形式
        #if self.cuda:
        	#images = images.cuda()
        #pr = self.net(images)[0]将图片传入网络中进行预测
        #pr = F.softmax(pr.permute(1,2,0),dim = -1).cup().numpy().argmax(axis=-1) 对图片最终结果取softmax，相当于获得每一个像素点获得每一个类的概率，再对概率取axis相当于获得每一个像素点所属于的类
        #pr = pr[int((self.model_image_size[0]-nh)//2)]:int((self.model_image_size[0]-nh)//2+nh),int((self.model_image_size[1]-nw)//2)对图像主体部分进行截取（去除灰条）
    #seg_img = np.zeros((np.shape(pr)[0],np.shape(pr)[1],3)) 对图像每个像素点所属种类进行判断，根据种类进行图像绘制
    #for c in range(self.num_classes):
    	#seg_img[:,:,0] += ((pr[:,: ] == c )*(self.colors[c][0])).astype('uint8')
        #seg_img[:,:,1] += ((pr[:,: ] == c )*(self.colors[c][1])).astype('uint8')
        #seg_img[:,:,2] += ((pr[:,: ] == c )*(self.colors[c][2])).astype('uint8')
    #image = Image.fromarray(np.uint8(seg_img)).resize((orininal_w,orininal_h))对图像进行resize
    #if self.blend:
    	#image = Image.blend(old_img,image,0.7)对原图进行混合
while True:
    img = input('Input image filename:')
    try:
        image = Image.open(img)
    except:
        print('Open Error! Try again!')
        continue
    else:
        r_image = pspnet.datect_image(image)
        r_image.show()

训练参数

if __name__ == "__main__":
    inputs_size = [473,473,3]#默认情况
    log_dir = "logs/"
    
    NUM_CLASSES = 21#需修改成自己要分类的类数+1 
    
    dice_loss = False
    #建议
    #种类少(几类)时，设置为True
    #种类多(十几类)时，如果batch_size比较大(10以上),那么设置为True
    #种类多(十几类)时，如果batch_size比较小(10一下)，那么设置为False
    
    #主干网络与训练权重使用
    #mobilenet和resnet50
    pretrained = False#如果要从头开始训练则将False更改为True，再将之后模型加载的地方注释掉
    backbone = "mobilenet"
    
    #是否使用辅助分支,会占用大量显存
    aux_branch = False
    
    #设置下采样的倍数
    #8和16，若为8则会占用大量显存
    downsample_factor = 16
    
    #cuda使用
    Cuda = True
    
    model = PSPNet(num_classes=NUM_CLASSES, backbone=backbone, downsample_factor=downsample_factor, pretrained=pretrained, aux_branch=aux_branch).train()
    
    
    
    model_path = "model_data/pspnet_mobilenetv2.pth"
    # 加快模型训练的效率
    print('Loading weights into state dict...')
    model_dict = model.state_dict()
    pretrained_dict = torch.load(model_path)
    pretrained_dict = { 
   k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) ==  np.shape(v)}
    model_dict.update(pretrained_dict)
    model.load_state_dict(model_dict)
    print('Finished!')
    
    
    
    if True:
        lr = 1e-4#学习率
        Init_Epoch = 0
        Interval_Epoch = 50#将主干特征提取网络冻结起来训练的代数
        Batch_size = 8#根据显存调整

        optimizer = optim.Adam(model.parameters(),lr)
        lr_scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=1,gamma=0.95)

        train_dataset = PSPnetDataset(train_lines, inputs_size, NUM_CLASSES, True)
        val_dataset = PSPnetDataset(val_lines, inputs_size, NUM_CLASSES, False)
        gen = DataLoader(train_dataset, shuffle=True, batch_size=Batch_size, num_workers=1, pin_memory=True,
                                drop_last=True, collate_fn=pspnet_dataset_collate)
        gen_val = DataLoader(val_dataset, shuffle=True, batch_size=Batch_size, num_workers=4,pin_memory=True, 
                                drop_last=True, collate_fn=pspnet_dataset_collate)

        epoch_size      = max(1, len(train_lines)//Batch_size)
        epoch_size_val  = max(1, len(val_lines)//Batch_size)

        for param in model.backbone.parameters():
            param.requires_grad = False

        for epoch in range(Init_Epoch,Interval_Epoch):
            fit_one_epoch(model,epoch,epoch_size,epoch_size_val,gen,gen_val,Interval_Epoch,Cuda,aux_branch)
            lr_scheduler.step()
    
    if True:#第二轮解冻训练
        lr = 1e-5#模型解冻后学习率调小点
        Interval_Epoch = 50
        Epoch = 100
        Batch_size = 4

        optimizer = optim.Adam(model.parameters(),lr)
        lr_scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=1,gamma=0.95)

        train_dataset = PSPnetDataset(train_lines, inputs_size, NUM_CLASSES, True)
        val_dataset = PSPnetDataset(val_lines, inputs_size, NUM_CLASSES, False)
        gen = DataLoader(train_dataset, shuffle=True, batch_size=Batch_size, num_workers=4, pin_memory=True,
                                drop_last=True, collate_fn=pspnet_dataset_collate)
        gen_val = DataLoader(val_dataset, shuffle=True, batch_size=Batch_size, num_workers=4,pin_memory=True, 
                                drop_last=True, collate_fn=pspnet_dataset_collate)

        epoch_size      = max(1, len(train_lines)//Batch_size)
        epoch_size_val  = max(1, len(val_lines)//Batch_size)

        for param in model.backbone.parameters():
            param.requires_grad = True

        for epoch in range(Interval_Epoch,Epoch):
            fit_one_epoch(model,epoch,epoch_size,epoch_size_val,gen,gen_val,Epoch,Cuda,aux_branch)
            lr_scheduler.step()

    

我们使用的训练文件采用VOC的格式。

语义分割模型训练的文件分为两部分。

原图就是普通的RGB图像，标签就是灰度图或者8位彩色图。

原图的shape为[height, width, 3]，标签的shape就是[height, width]，对于标签而言，每个像素点的内容是一个数字，比如0、1、2、3、4、5……，代表这个像素点所属的类别。

语义分割的工作就是对原始的图片的每一个像素点进行分类，所以通过预测结果中每个像素点属于每个类别的概率与标签对比，可以对网络进行训练。

2、LOSS解析

本文所使用的LOSS由两部分组成：

1、Cross Entropy Loss。

2、Dice Loss。

Cross Entropy Loss就是普通的交叉熵损失，当语义分割平台利用Softmax对像素点进行分类的时候，进行使用。

Dice loss将语义分割的评价指标作为Loss，Dice系数是一种集合相似度度量函数，通常用于计算两个样本的相似度，取值范围在[

0,1]。

计算公式如下：

$s=2\frac{|X\cap Y|}{|X|+|Y|}$

就是预测结果和真实结果的交乘上2，除上预测结果加上真实结果。其值在0-1之间。越大表示预测结果和真实结果重合度越大。所以Dice系数是越大越好。

如果作为LOSS的话是越小越好，所以使得Dice loss = 1 – Dice，就可以将Loss作为语义分割的损失了。

import torch
import torch.nn.functional as F  
import numpy as np
from torch import nn
from torch.autograd import Variable
from random import shuffle
from matplotlib.colors import rgb_to_hsv, hsv_to_rgb
from PIL import Image
import cv2

def CE_Loss(inputs, target, num_classes=21):
    n, c, h, w = inputs.size()
    nt, ht, wt = target.size()
    if h != ht and w != wt:
        inputs = F.interpolate(inputs, size=(ht, wt), mode="bilinear", align_corners=True)

    temp_inputs = inputs.transpose(1, 2).transpose(2, 3).contiguous().view(-1, c)
    temp_target = target.view(-1)

    CE_loss  = nn.NLLLoss(ignore_index=num_classes)(F.log_softmax(temp_inputs, dim = -1), temp_target)
    return CE_loss

def Dice_loss(inputs, target, beta=1, smooth = 1e-5):
    n, c, h, w = inputs.size()
    nt, ht, wt, ct = target.size()
    
    if h != ht and w != wt:
        inputs = F.interpolate(inputs, size=(ht, wt), mode="bilinear", align_corners=True)
    temp_inputs = torch.softmax(inputs.transpose(1, 2).transpose(2, 3).contiguous().view(n, -1, c),-1)
    temp_target = target.view(n, -1, ct)

    #--------------------------------------------#
    # 计算dice loss
    #--------------------------------------------#
    tp = torch.sum(temp_target[...,:-1] * temp_inputs, axis=[0,1])
    fp = torch.sum(temp_inputs                       , axis=[0,1]) - tp
    fn = torch.sum(temp_target[...,:-1]              , axis=[0,1]) - tp

    score = ((1 + beta ** 2) * tp + smooth) / ((1 + beta ** 2) * tp + beta ** 2 * fn + fp + smooth)
    dice_loss = 1 - torch.mean(score)
    return dice_loss


训练结果

-------#
    # 计算dice loss
    #--------------------------------------------#
    tp = torch.sum(temp_target[...,:-1] * temp_inputs, axis=[0,1])
    fp = torch.sum(temp_inputs                       , axis=[0,1]) - tp
    fn = torch.sum(temp_target[...,:-1]              , axis=[0,1]) - tp

    score = ((1 + beta ** 2) * tp + smooth) / ((1 + beta ** 2) * tp + beta ** 2 * fn + fp + smooth)
    dice_loss = 1 - torch.mean(score)
    return dice_loss

预测结果

今天的文章PSPNet_PSPNET自己的数据集「建议收藏」分享到此就结束了，感谢您的阅读。