名称：U-Net: Convolutional Networks for Biomedical Image Segmentation

网络结构

特点：U型网络 + Skip Connection跳层连接

左侧为下采样encoder，右侧为上采样decoder

灰色为Skip Connection：在上采样过程中融合下采样的feature map；将feature map通道叠加concat，concat可以通过crop或padding来匹配大小，UNet使用的是将小的feature map进行padding 0。

Double Conv模块

连续两次的卷积操作，每次3*3的unpadded卷积完接一个ReLU

使用BN操作来保持数据的分布

Down模块

下采样模块，有4次下采样，Double Conv模块后接2*2max pooling stride=2

Up模块

上采样模块，有4次上采样和特征融合的过程。

上采样过程包括2*2上卷积采样，特征融合，Double Conv，使用双线性插值和反卷积进行上采样。

特征融合过程先对小的feature map进行padding，然后concat

前向过程中，使用x1接收上采样数据，x2接收特征融合数据。

Out Conv模块

使用1*1卷积来映射结果类别

模型训练

加载数据

加载数据的伪代码如下

# ================================================================== #
#                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)

首先在__init__函数中指定目录和文件，然后在__getitem__函数中读取data，预处理data，返回data pair，___len__返回数据集大小。

然后创建Dataset类，通过DataLoader加载dataset，加载时选择batch_size和是否打乱。

修改网络

原网络的输出的尺寸会小于输入，因此对网络进行微调。

padding设为1，变成same padding

训练过程就是：加载数据集 => 定义优化算法 => 定义loss => 训练 => 记录最小的loss和模型参数

同时记得设置device、加载网络、拷贝数据和网络

预测

选择设备 => 加载网络结构 => 加载设备 => 加载模型 => 设置评估模式 => 设置图片路径

图片处理成单通道tensor => 获取预测结果=> 处理预测结果

Reference

https://cuijiahua.com/blog/2019/12/dl-15.html

https://cloud.tencent.com/developer/article/1633363

Mirror's live

UNet学习笔记