语义分割实战

语义分割基础与环境搭建

代码

语义分割基础与环境搭建

语义分割

语义分割（semantic segmentation） : 就是按照“语义”给图像上目标类别中的每一点打一个标签，使得不同种类的东西在图像上被区分开来。可以理解成像素级别的分类任务，直白点，就是对每个像素点进行分类。

简而言之，我们的目标是给定一幅RGB彩色图像（高x宽x3）或一幅灰度图像（高x宽x1），输出一个分割图谱，其中包括每个像素的类别标注（高x宽x1）。具体如下图所示：

注意：为了视觉上清晰，上面的预测图是一个低分辨率的图。在实际应用中，分割标注的分辨率需要与原始图像的分辨率相同。

这里对图片分为五类：Person（人）、Purse（包）、Plants/Grass（植物/草）、Sidewalk（人行道）、Building/Structures（建筑物）。

与标准分类值（standard categorical values）的做法相似，这里也是创建一个one-hot编码的目标类别标注——本质上即为每个类别创建一个输出通道。因为上图有5个类别，所以网络输出的通道数也为5，如下图所示：

如上图所示，预测的结果可以通过对每个像素在深度上求argmax的方式被整合到一张分割图中。进而，我们可以轻松地通过重叠的方式观察到每个目标。

argmax的方式也很好理解。如上图所示，每个通道只有0或1，以Person的通道为例，红色的1表示为Person的像素，其他像素均为0。其他通道也是如此，并且不存在同一个像素点在两个以上的通道均为1的情况。因此，通过argmax就找到每个像素点的最大索引通道值。最终得到结果为：

当只有一层通道被重叠至原始图像时，我们称之为mask，即只指示某一特定类别所存在的区域。

高分辨率的结果如下图所示，不同的颜色代表不同的类别：

数据集

第一个常用的数据集是Pascal VOC系列。这个系列中目前较流行的是VOC2012，Pascal Context等类似的数据集也有用到。

第二个常用的数据集是Microsoft COCO。COCO一共有80个类别，虽然有很详细的像素级别的标注，但是官方没有专门对语义分割的评测。这个数据集主要用于实例级别的分割以及图片描述。所以COCO数据集往往被当成是额外的训练数据集用于模型的训练。

第三个数据集是辅助驾驶（自动驾驶）环境的Cityscapes，使用比较常见的19个类别用于评测。

可以用于语义分割训练的数据集有很多：

Pascal Voc 2012：比较常见的物体分类，共21个类别；
MS COCO：由微软赞助，几乎成为了图像语义理解算法性能评价的“标准”数据集，共80个类别；
Cityscapes：包含50个欧洲城市不同场景、不同背景、不同季节的街景的33类标注物体；
Pascal-Context：对于PASCAL-VOC 2010识别竞赛的扩展，共59个类别；
KITTI：用于移动机器人及自动驾驶研究的最受欢迎的数据集之一，共11个类别；
NYUDv2：2.5维数据集，它包含1449张由微软Kinect设备捕获的室内的RGB-D图像；
SUN-RGBD：由四个RGB-D传感器得来，包含10000张RGB-D图像，尺寸与PASCAL VOC一致；
ADE20K_MIT：一个场景理解的新的数据集，这个数据集是可以免费下载的，共151个类别。

开发环境搭建

CUDA
Anaconda3
cuDNN和Pytorch安装

cuDNN是用于深度神经网络的GPU加速库。它强调性能、易用性和低内存开销。

UNet语义分割网络

代码

UNet论文

pytorch教程

UNet网络结构

在语义分割领域，基于深度学习的语义分割算法开山之作是FCN（Fully Convolutional Networks for Semantic Segmentation），而UNet是遵循FCN的原理，并进行了相应的改进，使其适应小样本的简单分割问题。

研究一个深度学习算法，可以先看网络结构，看懂网络结构后，再Loss计算方法、训练方法等。本文主要针对UNet的网络结构进行讲解，其它内容会在后续章节进行说明。

原理

UNet成为了大多做医疗影像语义分割任务的baseline。

UNet网络结构，最主要的两个特点是：U型网络结构和Skip Connection跳层连接。

UNet是一个对称的网络结构，左侧为下采样，右侧为上采样。
按照功能可以将左侧的一系列下采样操作称为encoder，将右侧的一系列上采样操作称为decoder。
Skip Connection中间四条灰色的平行线，Skip Connection就是在上采样的过程中，融合下采样过过程中的feature map。
Skip Connection用到的融合的操作也很简单，就是将feature map的通道进行叠加，俗称Concat。
Concat操作也很好理解，举个例子：一本大小为10cm*10cm，厚度为3cm的书A，和一本大小为10cm*10cm，厚度为4cm的书B。

将书A和书B，边缘对齐地摞在一起。这样就得到了，大小为10cm*10cm厚度为7cm的一摞书，类似这种“摞在一起”的操作，就是Concat。
同样道理，对于feature map，一个大小为256*256*64的feature map，即feature map的w（宽）为256，h（高）为256，c（通道数）为64。和一个大小为256*256*32的feature map进行Concat融合，就会得到一个大小为256*256*96的feature map。

在实际使用中，Concat融合的两个feature map的大小不一定相同，例如256*256*64的feature map和240*240*32的feature map进行Concat。

将大256*256*64的feature map进行裁剪，裁剪为240*240*64的feature map，比如上下左右，各舍弃8 pixel，裁剪后再进行Concat，得到240*240*96的feature map。
将小240*240*32的feature map进行padding操作，padding为256*256*32的feature map，比如上下左右，各补8 pixel，padding后再进行Concat，得到256*256*96的feature map。
UNet采用的Concat方案就是第二种，将小的feature map进行padding，padding的方式是补0，一种常规的常量填充。

代码

我们将整个UNet网络拆分为多个模块进行讲解。

DoubleConv模块：

连续两次的卷积操作。

从UNet网络中可以看出，不管是下采样过程还是上采样过程，每一层都会连续进行两次卷积操作，这种操作在UNet网络中重复很多次，可以单独写一个DoubleConv模块：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20


import torch.nn as nn

class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""
	# torch.nn.Sequential是一个时序容器，Module会以他们传入的顺序被添加到容器中
    # 卷积->BN->RELU->卷积->BN->RELU
    # DoubleConv模块的in_channels和out_channels可以灵活设定，以便扩展使用。
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=0),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)

如上图所示的网络，in_channels设为1，out_channels为64。输入图片大小为572*572，经过步长为1，padding为0的3*3卷积，得到570*570的feature map，再经过一次卷积得到568*568的feature map。

计算公式：O=(H−F+2×P)/S+1

H为输入feature map的大小，O为输出feature map的大小，F为卷积核的大小，P为padding的大小，S为步长。

Down模块：

UNet网络一共有4次下采样过程，模块化代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)

这里的代码很简单，就是一个maxpool池化层，进行下采样，然后接一个DoubleConv模块。

Up模块：

上采样过程用到的最多的当然就是上采样了，除了常规的上采样操作，还有进行特征的融合。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27


class Up(nn.Module):
    """Upscaling then double conv"""
    # 定义的上采样方法以及卷积采用DoubleConv。上采样，定义了两种方法：Upsample和ConvTranspose2d，也就是双线性插值和反卷积。
    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        else:
            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)

        self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        # input is CHW
        diffY = torch.tensor([x2.size()[2] - x1.size()[2]])
        diffX = torch.tensor([x2.size()[3] - x1.size()[3]])

        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])
        # if you have padding issues, see
        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)

上采样，定义了两种方法：Upsample和ConvTranspose2d，也就是双线性插值和反卷积。

双线性插值：

简单地讲：已知Q11、Q12、Q21、Q22四个点坐标，通过Q11和Q21求R1，再通过Q12和Q22求R2，最后通过R1和R2求P，这个过程就是双线性插值。

对于一个feature map而言，其实就是在像素点中间补点，补的点的值是多少，是由相邻像素点的值决定的。

反卷积，顾名思义，就是反着卷积。卷积是让featuer map越来越小，反卷积就是让feature map越来越大，示意图：

下面蓝色为原始图片，周围白色的虚线方块为padding结果，通常为0，上面绿色为卷积后的图片。

这个示意图，就是一个从2*2的feature map->4*4的feature map过程。

在forward前向传播函数中，x1接收的是上采样的数据，x2接收的是特征融合的数据。特征融合方法就是，上文提到的，先对小的feature map进行padding，再进行concat。

OutConv模块：

用上述的DoubleConv模块、Down模块、Up模块就可以拼出UNet的主体网络结构了。UNet网络的输出需要根据分割数量，整合输出通道，结果如下图所示：

至此，UNet网络用到的模块都已经写好，我们可以将上述的模块代码都放到一个unet_parts.py文件里，然后再创建unet_model.py，根据UNet网络结构，设置每个模块的输入输出通道个数以及调用顺序，编写如下代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40


import torch.nn.functional as F

from unet_parts import *


class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=False):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        self.down4 = Down(512, 1024)
        self.up1 = Up(1024, 512, bilinear)
        self.up2 = Up(512, 256, bilinear)
        self.up3 = Up(256, 128, bilinear)
        self.up4 = Up(128, 64, bilinear)
        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits


if __name__ == '__main__':
    net = UNet(n_channels=3, n_classes=1)
    print(net)

使用命令python unet_model.py，如果没有错误，你会得到如下结果：

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127


UNet(
  (inc): DoubleConv(
    (double_conv): Sequential(
      (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))
      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (2): ReLU(inplace=True)
      (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
      (4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (5): ReLU(inplace=True)
    )
  )
  (down1): Down(
    (maxpool_conv): Sequential(
      (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (1): DoubleConv(
        (double_conv): Sequential(
          (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1))
          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
          (3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
          (4): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (5): ReLU(inplace=True)
        )
      )
    )
  )
  (down2): Down(
    (maxpool_conv): Sequential(
      (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (1): DoubleConv(
        (double_conv): Sequential(
          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1))
          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
          (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
          (4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (5): ReLU(inplace=True)
        )
      )
    )
  )
  (down3): Down(
    (maxpool_conv): Sequential(
      (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (1): DoubleConv(
        (double_conv): Sequential(
          (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1))
          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
          (3): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1))
          (4): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (5): ReLU(inplace=True)
        )
      )
    )
  )
  (down4): Down(
    (maxpool_conv): Sequential(
      (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (1): DoubleConv(
        (double_conv): Sequential(
          (0): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1))
          (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
          (3): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1))
          (4): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (5): ReLU(inplace=True)
        )
      )
    )
  )
  (up1): Up(
    (up): ConvTranspose2d(1024, 512, kernel_size=(2, 2), stride=(2, 2))
    (conv): DoubleConv(
      (double_conv): Sequential(
        (0): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1))
        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): ReLU(inplace=True)
        (3): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1))
        (4): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (5): ReLU(inplace=True)
      )
    )
  )
  (up2): Up(
    (up): ConvTranspose2d(512, 256, kernel_size=(2, 2), stride=(2, 2))
    (conv): DoubleConv(
      (double_conv): Sequential(
        (0): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1))
        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): ReLU(inplace=True)
        (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
        (4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (5): ReLU(inplace=True)
      )
    )
  )
  (up3): Up(
    (up): ConvTranspose2d(256, 128, kernel_size=(2, 2), stride=(2, 2))
    (conv): DoubleConv(
      (double_conv): Sequential(
        (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1))
        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): ReLU(inplace=True)
        (3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
        (4): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (5): ReLU(inplace=True)
      )
    )
  )
  (up4): Up(
    (up): ConvTranspose2d(128, 64, kernel_size=(2, 2), stride=(2, 2))
    (conv): DoubleConv(
      (double_conv): Sequential(
        (0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1))
        (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): ReLU(inplace=True)
        (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
        (4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (5): ReLU(inplace=True)
      )
    )
  )
  (outc): OutConv(
    (conv): Conv2d(64, 1, kernel_size=(1, 1), stride=(1, 1))
  )
)

UNet模型训练

深度学习算法，无非就是我们解决一个问题的方法。选择什么样的网络去训练，进行什么样的预处理，采用什么Loss和优化方法，都是根据具体的任务而定的。

UNet 论文中的经典任务：医学图像分割。

简单描述一个这个任务：如动图所示，给一张细胞结构图，我们要把每个细胞互相分割开来。

这个训练数据只有30张，分辨率为512x512，这些图片是果蝇的电镜图。

UNet训练

想要训练一个深度学习模型，可以简单分为三个步骤：

数据加载：数据怎么加载，标签怎么定义，用什么数据增强方法，都是这一步进行。
模型选择：模型我们已经准备好了，就是该系列上篇文章讲到的 UNet 网络。
算法选择：算法选择也就是我们选什么 loss ，用什么优化算法。

每个步骤说的比较笼统，我们结合今天的医学图像分割任务，展开说明。

数据加载

这一步，可以做很多事情，说白了，无非就是图片怎么加载，标签怎么定义，为了增加算法的鲁棒性或者增加数据集，可以做一些数据增强的操作。

ISBI数据集

数据分为训练集和测试集，各30张，训练集有标签，测试集没有标签。

数据加载要做哪些处理，是根据任务和数据集而决定的，对于我们的分割任务，不用做太多处理，但由于数据量很少，仅30张，我们可以使用一些数据增强方法，来扩大我们的数据集。

Pytorch 给我们提供了一个方法，方便我们加载数据，我们可以使用这个框架，去加载我们的数据。看下伪代码：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25


# ================================================================== #
#                Input pipeline for custom dataset                 #
# ================================================================== #

# You should build your custom dataset as below.
class CustomDataset(torch.utils.data.Dataset):
    def __init__(self):
        # TODO
        # 1. Initialize file paths or a list of file names. 
        pass
    def __getitem__(self, index):
        # TODO
        # 1. Read one data from file (e.g. using numpy.fromfile, PIL.Image.open).
        # 2. Preprocess the data (e.g. torchvision.Transform).
        # 3. Return a data pair (e.g. image and label).
        pass
    def __len__(self):
        # You should change 0 to the total size of your dataset.
        return 0 

# You can then use the prebuilt data loader. 
custom_dataset = CustomDataset()
train_loader = torch.utils.data.DataLoader(dataset=custom_dataset,
                                           batch_size=64, 
                                           shuffle=True)

这是一个标准的模板，我们就使用这个模板，来加载数据，定义标签，以及进行数据增强。

创建一个dataset.py文件，编写代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55


import torch
import cv2
import os
import glob
from torch.utils.data import Dataset
import random


class ISBI_Loader(Dataset):
    def __init__(self, data_path):
        # 初始化函数，读取所有data_path下的图片
        self.data_path = data_path
        self.imgs_path = glob.glob(os.path.join(data_path, 'image/*.png'))

    def augment(self, image, flipCode):
        # 使用cv2.flip进行数据增强，filpCode为1水平翻转，0垂直翻转，-1水平+垂直翻转
        flip = cv2.flip(image, flipCode)
        return flip

    def __getitem__(self, index):
        # 根据index读取图片
        image_path = self.imgs_path[index]
        # 根据image_path生成label_path
        label_path = image_path.replace('image', 'label')
        # 读取训练图片和标签图片
        image = cv2.imread(image_path)
        label = cv2.imread(label_path)
        # 将数据转为单通道的图片
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        label = cv2.cvtColor(label, cv2.COLOR_BGR2GRAY)
        image = image.reshape(1, image.shape[0], image.shape[1])
        label = label.reshape(1, label.shape[0], label.shape[1])
        # 处理标签，将像素值为255的改为1
        if label.max() > 1:
            label = label / 255
        # 随机进行数据增强，为2时不做处理
        flipCode = random.choice([-1, 0, 1, 2])
        if flipCode != 2:
            image = self.augment(image, flipCode)
            label = self.augment(label, flipCode)
        return image, label

    def __len__(self):
        # 返回训练集大小
        return len(self.imgs_path)


if __name__ == "__main__":
    isbi_dataset = ISBI_Loader("data/train/")
    print("数据个数：", len(isbi_dataset))
    train_loader = torch.utils.data.DataLoader(dataset=isbi_dataset,
                                               batch_size=2,
                                               shuffle=True)
    for image, label in train_loader:
        print(image.shape)

运行结果：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16


数据个数： 30
torch.Size([2, 1, 512, 512])
torch.Size([2, 1, 512, 512])
torch.Size([2, 1, 512, 512])
torch.Size([2, 1, 512, 512])
torch.Size([2, 1, 512, 512])
torch.Size([2, 1, 512, 512])
torch.Size([2, 1, 512, 512])
torch.Size([2, 1, 512, 512])
torch.Size([2, 1, 512, 512])
torch.Size([2, 1, 512, 512])
torch.Size([2, 1, 512, 512])
torch.Size([2, 1, 512, 512])
torch.Size([2, 1, 512, 512])
torch.Size([2, 1, 512, 512])
torch.Size([2, 1, 512, 512])

__init__函数是这个类的初始化函数，根据指定的图片路径，读取所有图片数据，存放到self.imgs_path列表中。

__len__函数可以返回数据的多少，这个类实例化后，通过len()函数调用。

__getitem__函数是数据获取函数，在这个函数里你可以写数据怎么读，怎么处理，并且可以一些数据预处理、数据增强都可以在这里进行。我这里的处理很简单，只是将图片读取，并处理成单通道图片。同时，因为 label 的图片像素点是0和255，因此需要除以255，变成0和1。同时，随机进行了数据增强。

augment函数是定义的数据增强函数，怎么处理都行，我这里只是进行了简单的旋转操作。

在这个类中，你不用进行一些打乱数据集的操作，也不用管怎么按照 batchsize 读取数据。因为实例化这个类后，我们可以用 torch.utils.data.DataLoader 方法指定 batchsize 的大小，决定是否打乱数据。

Pytorch 提供给给我们的 DataLoader 很强大，我们甚至可以指定使用多少个进程加载数据，数据是否加载到 CUDA 内存中等高级用法，本文不涉及，就不再展开讲解了。

模型选择

但是我们需要对网络进行微调，完全按照论文的结构，模型输出的尺寸会稍微小于图片输入的尺寸，如果使用论文的网络结构需要在结果输出后，做一个 resize 操作。为了省去这一步，我们可以修改网络，使网络的输出尺寸正好等于图片的输入尺寸。

创建unet_parts.py文件

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67


import torch
import torch.nn as nn
import torch.nn.functional as F

class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)

class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)

class Up(nn.Module):
    """Upscaling then double conv"""

    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        else:
            self.up = nn.ConvTranspose2d(in_channels // 2, in_channels // 2, kernel_size=2, stride=2)

        self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        # input is CHW
        diffY = torch.tensor([x2.size()[2] - x1.size()[2]])
        diffX = torch.tensor([x2.size()[3] - x1.size()[3]])

        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2])

        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)

创建unet_model.py文件

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36


from model.unet_parts import *

class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        self.down4 = Down(512, 512)
        self.up1 = Up(1024, 256, bilinear)
        self.up2 = Up(512, 128, bilinear)
        self.up3 = Up(256, 64, bilinear)
        self.up4 = Up(128, 64, bilinear)
        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits

if __name__ == '__main__':
    net = UNet(n_channels=3, n_classes=1)
    print(net)

这样调整过后，网络的输出尺寸就与图片的输入尺寸相同了。

算法选择

选择什么 Loss 很重要，Loss 选择的好坏，都会影响算法拟合数据的效果。

选择什么 Loss 也是根据任务而决定的。我们今天的任务，只需要分割出细胞边缘，也就是一个很简单的二分类任务，所以我们可以使用 BCEWithLogitsLoss。

啥是 BCEWithLogitsLoss？BCEWithLogitsLoss 是 Pytorch 提供的用来计算二分类交叉熵的函数。

具体的公式推导，可以看我的机器学习系列教程《机器学习实战教程（六）：Logistic回归基础篇之梯度上升算法》，这里就不再累述。

目标函数，也就是 Loss 确定好了，怎么去优化这个目标呢？

最简单的方法就是，我们耳熟能详的梯度下降算法，逐渐逼近局部的极值。但是这种简单的优化算法，求解速度慢。各种优化算法，本质上其实都是梯度下降，例如最常规的 SGD，就是基于梯度下降改进的随机梯度下降算法，Momentum 就是引入了动量的 SGD，以指数衰减的形式累计历史梯度。

除了这些最基本的优化算法，还有自适应参数的优化算法。这类算法最大的特点就是，每个参数有不同的学习率，在整个学习过程中自动适应这些学习率，从而达到更好的收敛效果。

本文就是选择了一种自适应的优化算法 RMSProp。RMSProp 是基于 AdaGrad 的改进。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51


from model.unet_model import UNet
from utils.dataset import ISBI_Loader
from torch import optim
import torch.nn as nn
import torch

def train_net(net, device, data_path, epochs=40, batch_size=1, lr=0.00001):
    # 加载训练集
    isbi_dataset = ISBI_Loader(data_path)
    train_loader = torch.utils.data.DataLoader(dataset=isbi_dataset,
                                               batch_size=batch_size, 
                                               shuffle=True)
    # 定义RMSprop算法
    optimizer = optim.RMSprop(net.parameters(), lr=lr, weight_decay=1e-8, momentum=0.9)
    # 定义Loss算法
    criterion = nn.BCEWithLogitsLoss()
    # best_loss统计，初始化为正无穷
    best_loss = float('inf')
    # 训练epochs次
    for epoch in range(epochs):
        # 训练模式
        net.train()
        # 按照batch_size开始训练
        for image, label in train_loader:
            optimizer.zero_grad()
            # 将数据拷贝到device中
            image = image.to(device=device, dtype=torch.float32)
            label = label.to(device=device, dtype=torch.float32)
            # 使用网络参数，输出预测结果
            pred = net(image)
            # 计算loss
            loss = criterion(pred, label)
            print('Loss/train', loss.item())
            # 保存loss值最小的网络参数
            if loss < best_loss:
                best_loss = loss
                torch.save(net.state_dict(), 'best_model.pth')
            # 更新参数
            loss.backward()
            optimizer.step()

if __name__ == "__main__":
    # 选择设备，有cuda用cuda，没有就用cpu
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 加载网络，图片单通道1，分类为1。
    net = UNet(n_channels=1, n_classes=1)
    # 将网络拷贝到deivce中
    net.to(device=device)
    # 指定训练集地址，开始训练
    data_path = "data/train/"
    train_net(net, device, data_path)

预测

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43


import glob
import numpy as np
import torch
import os
import cv2
from model.unet_model import UNet

if __name__ == "__main__":
    # 选择设备，有cuda用cuda，没有就用cpu
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 加载网络，图片单通道，分类为1。
    net = UNet(n_channels=1, n_classes=1)
    # 将网络拷贝到deivce中
    net.to(device=device)
    # 加载模型参数
    net.load_state_dict(torch.load('best_model.pth', map_location=device))
    # 测试模式
    net.eval()
    # 读取所有图片路径
    tests_path = glob.glob('data/test/*.png')
    # 遍历所有图片
    for test_path in tests_path:
        # 保存结果地址
        save_res_path = test_path.split('.')[0] + '_res.png'
        # 读取图片
        img = cv2.imread(test_path)
        # 转为灰度图
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        # 转为batch为1，通道为1，大小为512*512的数组
        img = img.reshape(1, 1, img.shape[0], img.shape[1])
        # 转为tensor
        img_tensor = torch.from_numpy(img)
        # 将tensor拷贝到device中，只用cpu就是拷贝到cpu中，用cuda就是拷贝到cuda中。
        img_tensor = img_tensor.to(device=device, dtype=torch.float32)
        # 预测
        pred = net(img_tensor)
        # 提取结果
        pred = np.array(pred.data.cpu()[0])[0]
        # 处理结果
        pred[pred >= 0.5] = 255
        pred[pred < 0.5] = 0
        # 保存图片
        cv2.imwrite(save_res_path, pred)

查看loss

https://cuijiahua.com/blog/2020/05/dl-17.html

训练模型，最常看的指标就是 Loss。我们可以根据 Loss 的收敛情况，初步判断模型训练的好坏。

如果，Loss 值突然上升了，那说明训练有问题，需要检查数据和代码。

如果，Loss 值趋于稳定，那说明训练完毕了。

观察 Loss 情况，最直观的方法，就是绘制 Loss 曲线图。

sys.stdout

通过 Loss 曲线，我们可以分析模型训练的好坏，模型是否训练完成，起到一个很好的“监控”作用。

绘制 Loss 曲线图，第一步就是需要保存训练过程中的 Loss 值。

一个最简单的方法是使用，sys.stdout 标准输出重定向，简单好用。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


import os
import sys
class Logger():
    def __init__(self, filename="log.txt"):
        self.terminal = sys.stdout
        self.log = open(filename, "w")

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)

    def flush(self):
        pass

sys.stdout = Logger()

print("123456")
print("https://cuijiahua.com")

代码很简单，创建一个 log.py 文件，自己写一个 Logger 类，并采用 sys.stdout 重定向输出。

在 Terminal 中，不仅可以使用 print 打印结果，同时也会将结果保存到 log.txt 文件中。

运行 log.py，打印 print 内容的同时，也将内容写入了 log.txt 文件中。

使用这个代码，就可以在打印 Loss 的同时，将结果保存到指定的 txt 中，比如保存上篇文章训练 UNet 的 Loss。

1
2
3
4


from log import *

if __name__ == "__main__":
    sys.stdout = Logger()

matplotlib

Matplotlib 是一个 Python 的绘图库，简单好用。

简单几行命令，就可以绘制曲线图、散点图、条形图、直方图、饼图等等。

在深度学习中，一般就是绘制曲线图，比如 Loss 曲线、Acc 曲线。

使用 sys.stdout 保存的 train_loss.txt，绘制 Loss 曲线。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


import matplotlib.pyplot as plt
with open('train_loss.txt', 'r') as f:
    train_loss = f.readlines()
    train_loss = list(map(lambda x:float(x.strip()), train_loss))
x = range(len(train_loss))
y = train_loss
plt.plot(x, y, label='train loss', linewidth=2, color='r', marker='o', markerfacecolor='r', markersize=5)
plt.xlabel('Epoch')
plt.ylabel('Loss Value')
plt.legend()
plt.show()

Logging

说到保存日志，那不得不提 Python 的内置标准模块 Logging，它主要用于输出运行日志，可以设置输出日志的等级、日志保存路径、日志文件回滚等，同时，我们也可以设置日志的输出格式。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


import logging

def get_logger(LEVEL, log_file = None):
    head = '[%(asctime)-15s] [%(levelname)s] %(message)s'
    if LEVEL == 'info':
        logging.basicConfig(level=logging.INFO, format=head)
    elif LEVEL == 'debug':
        logging.basicConfig(level=logging.DEBUG, format=head)
    logger = logging.getLogger()
    if log_file != None:
        fh = logging.FileHandler(log_file)
        logger.addHandler(fh)
    return logger

logger = get_logger('info')

logger.info('123456')
logger.info('https://cuijiahua.com')

只需要几行代码，进行一个简单的封装使用。使用函数 get_logger 创建一个级别为 info 的 logger，如果指定 log_file，则会对日志进行保存。

当然，我们可以通过，logging.basicConfig 的 format 参数，设置日志格式。

TensorboardX

TensorboardX ，它是专门用于深度学习“炼丹”的高级“法宝”。

在 Pytorch 中，这个可视化工具叫做 TensorBoardX，其实就是针对 Tensorboard 的一个封装，使得 PyTorch 用户也能够调用 Tensorboard。

TensorboardX 安装也非常简单，使用 pip 即可安装，需要注意的是 Pytorch 的版本需要大于 1.1.0。

1

pip install tensorboardX

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


from tensorboardX import SummaryWriter

# 创建 writer1 对象
# log 会保存到 runs/exp 文件夹中
writer1 = SummaryWriter('runs/exp')

# 使用默认参数创建 writer2 对象
# 自动创建runs文件夹
# log 会保存到 runs/日期_用户名 格式的文件夹中
writer2 = SummaryWriter()

# 使用 commet 参数，创建 writer3 对象
# log 会保存到 runs/日期_用户名_resnet 格式的文件中
writer3 = SummaryWriter(comment='_resnet')

使用的时候，创建一个 SummaryWriter 对象即可，以上展示了三种初始化 SummaryWriter 的方法：

提供一个路径，将使用该路径来保存日志
无参数，默认将使用 runs/日期_用户名路径来保存日志
提供一个 comment 参数，将使用 runs/日期_用户名+comment 路径来保存日志

运行结果：

有了 writer 我们就可以往日志里写入数字、图片、甚至声音等数据。

数字(scalar)

这个是最简单的，使用 add_scalar 方法来记录数字常量。

1

add_scalar(tag, scalar_value, global_step=None, walltime=None)

总共 4 个参数。

tag (string): 数据名称，不同名称的数据使用不同曲线展示
scalar_value (float): 数字常量值
global_step (int, optional): 训练的 step
walltime (float, optional): 记录发生的时间，默认为 time.time()

需要注意，这里的 scalar_value 一定是 float 类型，如果是 PyTorch scalar tensor，则需要调用 .item() 方法获取其数值。我们一般会使用 add_scalar 方法来记录训练过程的 loss、accuracy、learning rate 等数值的变化，直观地监控训练过程。

1
2
3
4
5
6


from tensorboardX import SummaryWriter    
writer = SummaryWriter('runs/scalar_example')
for i in range(10):
    writer.add_scalar('quadratic', i**2, global_step=i)
    writer.add_scalar('exponential', 2**i, global_step=i)
writer.close()

通过 add_scalar 往日志里写入数字，日志保存到 runs/scalar_example中，writer 用完要记得 close，否则无法保存数据。

在 cmd 中使用如下命令：

1

tensorboard --logdir=runs/scalar_example --port=8088

指定日志地址，使用端口号，在浏览器中，就可以使用如下地址，打开 Tensorboad。

1

http://localhost:8088/

不…太…行…

Contents