研二了,第一次用深度学习做分割,感觉自己落伍了好多,方向是图像处理,却一直在用传统法拼拼凑凑,同学都说深度学习要发文章得有好的数学基础,自知数学基础差的情况下还是要接触一下的,毕竟万事开头难,不学习就永远不会,那么闲话不多说,记录一下自己的学习过程。
(环境=py3.7+pytorch+spyder)
首先根据前人经验先搭建网络,下图是Unet的网络结构图:
观察到conv操作蛮多的,不管是下采样层还是反卷积层中都用到,那么先写一个该操作的class打包一下:
class Conv3x3(nn.Module):
def __init__(self, inputCh, outputCh):
super(Conv3x3, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(inputCh, outputCh, kernel_size=3, stride=1, padding=1),#卷积核3x3,in->out
nn.BatchNorm2d(outPutCh),#规范化
nn.ReLU(inplace=True),#激活函数ReLU
)
self.conv2 = nn.Sequential(
nn.Conv2d(outputCh, outputCh, kernel_size=3, stride=1, padding=1),#根据图,上一次的out->out
nn.BatchNorm2d(outputCh),
nn.ReLU(inplace=True),
)
def forward(self, x):#前向传播
x = self.conv1(x)
x = self.conv2(x)
return x
打包完卷积的操作之后,再把上采样的操作整理一下:
class TransConv(nn.Module):
def __init__(self, inputCh, outputCh):
super(TransConv, self).__init__()
self.conv = nn.Sequential(
nn.ConvTranspose2d(inputCh, outputCh, kernel_size=3, stride=2, padding=1, output_padding=1, dilation=1),
nn.BatchNorm2d(outputCh),
nn.ReLU(inplace=True),
)
def forward(self, x):
x = self.conv(x)
return x
class UpSam(nn.Module):
def __init__(self, inputCh, outputCh):
super(UpSam, self).__init__()
self.upconv = TransConv(inputCh, outputCh)#反卷积
self.conv = Conv3x3(2 * outputCh, outputCh)#这里用到上面写的conv操作
def forward(self, x, convfeatures):
x = self.upconv(x)
x = torch.cat([x, convfeatures], dim=1)
x = self.conv(x)
return x
至此完成图中蓝色箭头,灰色箭头,绿色箭头的定义,红色箭头是maxpool,实质是下采样,可以跟其他block组合到一起,整体网络如下:
class UNet(nn.Module):
def __init__(self, inputCh=4, outputCh=5, size=64):#4种模态数据,拟输出5个类别(label数据0~4表示:背景、坏死组织、囊肿、肿瘤核心、整体肿瘤)
super(UNet, self).__init__()
channels = []
for i in range(5):
channels.append((2 ** i) * size)#对应图像的size
self.downLayer1 = Conv3x3(inputCh, channels[0])
self.downLayer2 = nn.Sequential(nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
Conv3x3(channels[0], channels[1]))
self.downLayer3 = nn.Sequential(nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
Conv3x3(channels[1], channels[2]))
self.downLayer4 = nn.Sequential(nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
Conv3x3(channels[2], channels[3]))
self.bottomLayer = nn.Sequential(nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
Conv3x3(channels[3], channels[4]))
self.upLayer1 = UpSam(channels[4], channels[3])
self.upLayer2 = UpSam(channels[3], channels[2])
self.upLayer3 = UpSam(channels[2], channels[1])
self.upLayer4 = UpSam(channels[1], channels[0])
self.outLayer = nn.Conv2d(channels[0], outputCh, kernel_size=3, stride=1, padding=1)
def forward(self, x):
#前半条路
x1 = self.downLayer1(x) # size(32) * 16 * W * H
x2 = self.downLayer2(x1) # size(64) * 16/2 * W/2 * H/2
x3 = self.downLayer3(x2) # size(128) * 16/4 * W/4 * H/4
x4 = self.downLayer4(x3) # size(256) * 16/8 * W/8 * H/8
#最底层
x5 = self.bottomLayer(x4) # size(512) * 16/16 * W/16 * H/16
#后半条路
x = self.upLayer1(x5, x4) # size(256) * 16/8 * W/8 * H/8
x = self.upLayer2(x, x3) # size(128) * 16/4 * W/4 * H/4
x = self.upLayer3(x, x2) # size(64) * 16/2 * W/2 * H/2
x = self.upLayer4(x, x1) # size(32) * 16 * W * H
x = self.outLayer(x) # outputCh(2 ) * 16 * W * H
return x
网络构建完毕,写个main函数验证下看看:
if __name__ == "__main__":
net = UNet(4, 5, degree=64)
batch_size = 4
a = torch.randn(batch_size, 4, 192, 192)#随便搞点数据扔进去
b = net(a)
print(a.shape)
print(b.shape)
可以从调试结果看到,网络输出的结果与网络的输入是同维度的(4,192,192),输入包含4个模态,输出包含了5个类别,这与我们期望的结果吻合,应该没啥事问题,那么就准备写DataLoader了;
个人感觉写DataLoader最重要的是记得自己要输入什么样的数据,输出怎么样的矩阵,现在还在写,读取文件有点麻烦,改天再更新。
import sys
sys.path.append("..")
import math
import torch
import torch.nn as nn
class ConvBlock2d(nn.Module):
def __init__(self, in_ch, out_ch):
super(ConvBlock2d, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
)
self.conv2 = nn.Sequential(
nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
return x
class ConvTrans2d(nn.Module):
def __init__(self, in_ch, out_ch):
super(ConvTrans2d, self).__init__()
self.conv1 = nn.Sequential(
nn.ConvTranspose2d(in_ch, out_ch, kernel_size=3, stride=2, padding=1, output_padding=1, dilation=1),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
)
def forward(self, x):
x = self.conv1(x)
return x
class UpBlock2d(nn.Module):
def __init__(self, in_ch, out_ch):
super(UpBlock2d, self).__init__()
self.up_conv = ConvTrans2d(in_ch, out_ch)
self.conv = ConvBlock2d(2 * out_ch, out_ch)
def forward(self, x, down_features):
x = self.up_conv(x)
x = torch.cat([x, down_features], dim=1)
x = self.conv(x)
return x
class UNet2D(nn.Module):
def __init__(self, in_ch=4, out_ch=2, degree=64):
super(UNet2D, self).__init__()
chs = []
for i in range(5):
chs.append((2 ** i) * degree)
self.downLayer1 = ConvBlock2d(in_ch, chs[0])
self.downLayer2 = nn.Sequential(nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
ConvBlock2d(chs[0], chs[1]))
self.downLayer3 = nn.Sequential(nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
ConvBlock2d(chs[1], chs[2]))
self.downLayer4 = nn.Sequential(nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
ConvBlock2d(chs[2], chs[3]))
self.bottomLayer = nn.Sequential(nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
ConvBlock2d(chs[3], chs[4]))
self.upLayer1 = UpBlock2d(chs[4], chs[3])
self.upLayer2 = UpBlock2d(chs[3], chs[2])
self.upLayer3 = UpBlock2d(chs[2], chs[1])
self.upLayer4 = UpBlock2d(chs[1], chs[0])
self.outLayer = nn.Conv2d(chs[0], out_ch, kernel_size=3, stride=1, padding=1)
def forward(self, x):
x1 = self.downLayer1(x)
x2 = self.downLayer2(x1)
x3 = self.downLayer3(x2)
x4 = self.downLayer4(x3)
x5 = self.bottomLayer(x4)
x = self.upLayer1(x5, x4)
x = self.upLayer2(x, x3)
x = self.upLayer3(x, x2)
x = self.upLayer4(x, x1)
x = self.outLayer(x)
return x
if __name__ == "__main__":
net = UNet2D(4, 5, degree=64)
print("total parameter:" + str(netSize(net)))
batch_size = 4
a = torch.randn(batch_size, 4, 192, 192)
b = net(a)
print(b.shape)