憨批的語義分割重制版6——Pytorch 搭建自己的Unet語義分割平台

• 注意事項
• 學習前言
• 什麼是Unet模型
• 代碼下載下傳
• Unet實作思路
• • 一、預測部分
  • • 1、主幹網絡介紹
    • 2、加強特征提取結構
    • 3、利用特征獲得預測結果
  • 二、訓練部分
  • • 1、訓練檔案詳解
    • 2、LOSS解析
• 訓練自己的Unet模型

注意事項

這是重新建構了的Unet語義分割網絡，主要是檔案架構上的建構，還有代碼的實作，和之前的語義分割網絡相比，更加完整也更清晰一些。建議還是學習這個版本的Unet。

學習前言

還是快樂的pytorch人。

什麼是Unet模型

Unet是一個優秀的語義分割模型，其主要執行過程與其它語義分割模型類似。

Unet可以分為三個部分，如下圖所示：

第一部分是主幹特征提取部分，我們可以利用主幹部分獲得一個又一個的特征層，Unet的主幹特征提取部分與VGG相似，為卷積和最大池化的堆疊。利用主幹特征提取部分我們可以獲得五個初步有效特征層，在第二步中，我們會利用這五個有效特征層可以進行特征融合。

第二部分是加強特征提取部分，我們可以利用主幹部分擷取到的五個初步有效特征層進行上采樣，并且進行特征融合，獲得一個最終的，融合了所有特征的有效特征層。

第三部分是預測部分，我們會利用最終獲得的最後一個有效特征層對每一個特征點進行分類，相當于對每一個像素點進行分類。

代碼下載下傳

Github源碼下載下傳位址為：

https://github.com/bubbliiiing/unet-pytorch

Unet實作思路

一、預測部分

1、主幹網絡介紹

Unet的主幹特征提取部分由卷積+最大池化組成，整體結構與VGG類似。

本文所采用的主幹特征提取網絡為VGG16，這樣也友善使用imagnet上的預訓練權重。

VGG是由Simonyan 和Zisserman在文獻《Very Deep Convolutional Networks for Large Scale Image Recognition》中提出卷積神經網絡模型，其名稱來源于作者所在的牛津大學視覺幾何組(Visual Geometry Group)的縮寫。

該模型參加2014年的 ImageNet圖像分類與定位挑戰賽，取得了優異成績：在分類任務上排名第二，在定位任務上排名第一。

它的結構如下圖所示：

這是一個VGG16被用到爛的圖，但确實很好的反應了VGG16的結構。

當我們使用VGG16作為主幹特征提取網絡的時候，我們隻會用到兩種類型的層，分别是卷積層和最大池化層。

當輸入的圖像大小為512x512x3的時候，具體執行方式如下：

1、conv1：進行兩次[3,3]的64通道的卷積，獲得一個[512,512,64]的初步有效特征層，再進行2X2最大池化，獲得一個[256,256,64]的特征層。

2、conv2：進行兩次[3,3]的128通道的卷積，獲得一個[256,256,128]的初步有效特征層，再進行2X2最大池化，獲得一個[128,128,128]的特征層。

3、conv3：進行三次[3,3]的256通道的卷積，獲得一個[128,128,256]的初步有效特征層，再進行2X2最大池化，獲得一個[64,64,256]的特征層。

4、conv4：進行三次[3,3]的512通道的卷積，獲得一個[64,64,512]的初步有效特征層，再進行2X2最大池化，獲得一個[32,32,512]的特征層。

5、conv5：進行三次[3,3]的512通道的卷積，獲得一個[32,32,512]的初步有效特征層。

import torch
import torch.nn as nn
from torchvision.models.utils import load_state_dict_from_url

class VGG(nn.Module):
    def __init__(self, features, num_classes=1000):
        super(VGG, self).__init__()
        self.features = features
        self.avgpool = nn.AdaptiveAvgPool2d((7, 7))
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, num_classes),
        )
        self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)


def make_layers(cfg, batch_norm=False, in_channels = 3):
    layers = []
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)

cfgs = {
    'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']
}


def VGG16(pretrained, in_channels, **kwargs):
    model = VGG(make_layers(cfgs["D"], batch_norm = False, in_channels = in_channels), **kwargs)
    if pretrained:
        state_dict = load_state_dict_from_url("https://download.pytorch.org/models/vgg16-397923af.pth", model_dir="./model_data")
        model.load_state_dict(state_dict)
    
    del model.avgpool
    del model.classifier
    return model
           

2、加強特征提取結構

Unet所使用的加強特征提取網絡是一個U的形狀。

利用第一步我們可以獲得五個初步的有效特征層，在加強特征提取網絡這裡，我們會利用這五個初步的有效特征層進行特征融合，特征融合的方式就是對特征層進行上采樣并且進行堆疊。

為了友善網絡的建構與更好的通用性，我們的Unet和上圖的Unet結構有些許不同，在上采樣時直接進行兩倍上采樣再進行特征融合，最終獲得的特征層和輸入圖檔的高寬相同。

具體示意圖如下：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary
from nets.vgg import VGG16


class unetUp(nn.Module):
    def __init__(self, in_size, out_size):
        super(unetUp, self).__init__()
        self.conv1 = nn.Conv2d(in_size, out_size, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(out_size, out_size, kernel_size=3, padding=1)
        self.up = nn.UpsamplingBilinear2d(scale_factor=2)

    def forward(self, inputs1, inputs2):

        outputs = torch.cat([inputs1, self.up(inputs2)], 1)
        outputs = self.conv1(outputs)
        outputs = self.conv2(outputs)
        return outputs


class Unet(nn.Module):
    def __init__(self, num_classes=21, in_channels=3, pretrained=False):
        super(Unet, self).__init__()
        self.vgg = VGG16(pretrained=pretrained,in_channels=in_channels)
        in_filters = [192, 384, 768, 1024]
        out_filters = [64, 128, 256, 512]
        # upsampling
        self.up_concat4 = unetUp(in_filters[3], out_filters[3])
        self.up_concat3 = unetUp(in_filters[2], out_filters[2])
        self.up_concat2 = unetUp(in_filters[1], out_filters[1])
        self.up_concat1 = unetUp(in_filters[0], out_filters[0])

        # final conv (without any concat)
        self.final = nn.Conv2d(out_filters[0], num_classes, 1)

    def forward(self, inputs):
        feat1 = self.vgg.features[  :4 ](inputs)
        feat2 = self.vgg.features[4 :9 ](feat1)
        feat3 = self.vgg.features[9 :16](feat2)
        feat4 = self.vgg.features[16:23](feat3)
        feat5 = self.vgg.features[23:-1](feat4)

        up4 = self.up_concat4(feat4, feat5)
        up3 = self.up_concat3(feat3, up4)
        up2 = self.up_concat2(feat2, up3)
        up1 = self.up_concat1(feat1, up2)

        final = self.final(up1)
        
        return final

    def _initialize_weights(self, *stages):
        for modules in stages:
            for module in modules.modules():
                if isinstance(module, nn.Conv2d):
                    nn.init.kaiming_normal_(module.weight)
                    if module.bias is not None:
                        module.bias.data.zero_()
                elif isinstance(module, nn.BatchNorm2d):
                    module.weight.data.fill_(1)
                    module.bias.data.zero_()
           

3、利用特征獲得預測結果

利用1、2步，我們可以擷取輸入進來的圖檔的特征，此時，我們需要利用特征獲得預測結果。

利用特征獲得預測結果的過程為：

利用一個1x1卷積進行通道調整，将最終特征層的通道數調整成num_classes。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary
from nets.vgg import VGG16


class unetUp(nn.Module):
    def __init__(self, in_size, out_size):
        super(unetUp, self).__init__()
        self.conv1 = nn.Conv2d(in_size, out_size, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(out_size, out_size, kernel_size=3, padding=1)
        self.up = nn.UpsamplingBilinear2d(scale_factor=2)

    def forward(self, inputs1, inputs2):

        outputs = torch.cat([inputs1, self.up(inputs2)], 1)
        outputs = self.conv1(outputs)
        outputs = self.conv2(outputs)
        return outputs


class Unet(nn.Module):
    def __init__(self, num_classes=21, in_channels=3, pretrained=False):
        super(Unet, self).__init__()
        self.vgg = VGG16(pretrained=pretrained,in_channels=in_channels)
        in_filters = [192, 384, 768, 1024]
        out_filters = [64, 128, 256, 512]
        # upsampling
        self.up_concat4 = unetUp(in_filters[3], out_filters[3])
        self.up_concat3 = unetUp(in_filters[2], out_filters[2])
        self.up_concat2 = unetUp(in_filters[1], out_filters[1])
        self.up_concat1 = unetUp(in_filters[0], out_filters[0])

        # final conv (without any concat)
        self.final = nn.Conv2d(out_filters[0], num_classes, 1)

    def forward(self, inputs):
        feat1 = self.vgg.features[  :4 ](inputs)
        feat2 = self.vgg.features[4 :9 ](feat1)
        feat3 = self.vgg.features[9 :16](feat2)
        feat4 = self.vgg.features[16:23](feat3)
        feat5 = self.vgg.features[23:-1](feat4)

        up4 = self.up_concat4(feat4, feat5)
        up3 = self.up_concat3(feat3, up4)
        up2 = self.up_concat2(feat2, up3)
        up1 = self.up_concat1(feat1, up2)

        final = self.final(up1)
        
        return final

    def _initialize_weights(self, *stages):
        for modules in stages:
            for module in modules.modules():
                if isinstance(module, nn.Conv2d):
                    nn.init.kaiming_normal_(module.weight)
                    if module.bias is not None:
                        module.bias.data.zero_()
                elif isinstance(module, nn.BatchNorm2d):
                    module.weight.data.fill_(1)
                    module.bias.data.zero_()
           

二、訓練部分

1、訓練檔案詳解

我們使用的訓練檔案采用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]。

計算公式如下：

就是預測結果和真實結果的交乘上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
           

訓練自己的Unet模型

整個Unet的檔案構架為：

在訓練模型之前，我們需要首先準備好資料集。

大家可以下載下傳我上傳的voc資料集，也可以根據voc資料集格式進行資料集制作。

如果大家下載下傳的是我上傳的voc資料集，那麼就不需要運作VOCdevkit檔案夾下面的voc2unet.py。

如果是自己制作的資料集，那麼需要運作VOCdevkit檔案夾下面的voc2unet.py，進而生成train.txt和val.txt。

生成完成後。

修改train.py檔案下的num_classes，使其為分的類的個數+1。

之後就可以開始訓練了。