基于U-net的医学影像分割 / pytorch实现

一、简介

U-net是一种基于CNN神经网络的图像分割模型，模型主要包括两个部分：对图片的特征提取与上采样。模型的巧妙之处在于在上采样时，会将上采样结果与特征提取中对应通道的特征进行拼接，然后接着进行卷积运算。这样能更多地利用原图片的特征信息。网络结构图与大写字母U非常类似，因此称为U-net，模型结构图如下。

 

 

二、U-net实现

由于数图课程的作业要求，作为一个初学者，近期学习并实现了U-net模型，现将过程做一个简单的回顾，希望能给有需要的初学者一点小小的帮助。

1、实验环境

本次实验环境如下：Ubuntu 16.04，python 3.6，cuda 9.0，pytorch

数据集：数据使用ISBI Challenge 2012的数据（已存放在GitHub中）

2、U-net模型

import torch
from torch import nn

# 卷积层代码块，进行两次卷积晕眩
class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

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

# 模型实现，可参照模型图
class Unet(nn.Module):
    def __init__(self,in_ch,out_ch):
        super(Unet, self).__init__()

        self.conv1 = DoubleConv(in_ch, 64)
        self.pool1 = nn.MaxPool2d(2)
        self.conv2 = DoubleConv(64, 128)
        self.pool2 = nn.MaxPool2d(2)
        self.conv3 = DoubleConv(128, 256)
        self.pool3 = nn.MaxPool2d(2)
        self.conv4 = DoubleConv(256, 512)
        self.pool4 = nn.MaxPool2d(2)
        self.conv5 = DoubleConv(512, 1024)
        self.up6 = nn.ConvTranspose2d(1024, 512, 2, stride=2)
        self.conv6 = DoubleConv(1024, 512)
        self.up7 = nn.ConvTranspose