U-Net网络学习--01

摘要

本周主要进行u-net相关的学习，U-Net是比较早的使用全卷积网络进行语义分割的算法之一，该网络的名字也是取自其U形形状。UNet是一个经典的网络设计方式，在医学图像分割任务中具有大量的应用，其采用的编码器（下采样）-解码器（上采样）结构和跳跃连接是一种非常经典的设计方法。训练一个u-net的模型并及进行分析。

一、U-Net: Convolutional Networks for Biomedical Image Segmentation

原论文地址：https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/

1.1 U-Net的网络结构

U-Net的U形结构如图1所示。网络是一个经典的全卷积网络（即网络中没有全连接操作）。网络的输入是一张 [公式] 的边缘经过镜像操作的图片（input image tile），关于“镜像操作“会在1.2节进行详细分析，网络的左侧（红色虚线）是由卷积和Max Pooling构成的一系列降采样操作，论文中将这一部分叫做压缩路径（contracting path）。压缩路径由4个block组成，每个block使用了3个有效卷积和1个Max Pooling降采样，每次降采样之后Feature Map的个数乘2，因此有了图中所示的Feature Map尺寸变化。最终得到了尺寸为 [公式] 的Feature Map。

网络的右侧部分(绿色虚线)在论文中叫做扩展路径（expansive path）。同样由4个block组成，每个block开始之前通过反卷积将Feature Map的尺寸乘2，同时将其个数减半（最后一层略有不同），然后和左侧对称的压缩路径的Feature Map合并，由于左侧压缩路径和右侧扩展路径的Feature Map的尺寸不一样，U-Net是通过将压缩路径的Feature Map裁剪到和扩展路径相同尺寸的Feature Map进行归一化的（即图1中左侧虚线部分）。扩展路径的卷积操作依旧使用的是有效卷积操作，最终得到的Feature Map的尺寸是 [公式] 。由于该任务是一个二分类任务，所以网络有两个输出Feature Map。

1.2 U-Net究竟输入了什么

首先，数据集我们的原始图像的尺寸都是 [公式] 的。为了能更好的处理图像的边界像素，U-Net使用了镜像操作（Overlay-tile Strategy）来解决该问题。镜像操作即是给输入图像加入一个对称的边（图2），那么边的宽度是多少呢？一个比较好的策略是通过感受野确定。因为有效卷积是会降低Feature Map分辨率的，但是我们希望 [公式] 的图像的边界点能够保留到最后一层Feature Map。所以我们需要通过加边的操作增加图像的分辨率，增加的尺寸即是感受野的大小，也就是说每条边界增加感受野的一半作为镜像边。

1.3 U-Net的损失函数

ISBI数据集的一个非常严峻的挑战是紧密相邻的物体之间的分割问题。如图3所示，(a)是输入数据，(b)是Ground Truth，©是基于Ground Truth生成的分割掩码，(d)是U-Net使用的用于分离边界的损失权值。

二、deep-learning-for-image-processing

下面是原作者的网址，可以去下载压缩包，根据提示配环境并运行测试：https://github.com/WZMIAOMIAO/deep-learning-for-image-processing

2.1 训练代码

import os
import time
import datetime

import torch

from src import UNet
from train_utils import train_one_epoch, evaluate, create_lr_scheduler
from my_dataset import DriveDataset
import transforms as T


class SegmentationPresetTrain:
    def __init__(self, base_size, crop_size, hflip_prob=0.5, vflip_prob=0.5,
                 mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
        min_size = int(0.5 * base_size)
        max_size = int(1.2 * base_size)

        trans = [T.RandomResize(min_size, max_size)]
        if hflip_prob > 0:
            trans.append(T.RandomHorizontalFlip(hflip_prob))
        if vflip_prob > 0:
            trans.append(T.RandomVerticalFlip(vflip_prob))
        trans.extend([
            T.RandomCrop(crop_size),
            T.ToTensor(),
            T.Normalize(mean=mean, std=std),
        ])
        self.transforms = T.Compose(trans)

    def __call__(self, img, target):
        return self