unet网络python代码详解_使用pytorch实现论文中的unet网络
设计神经网络的一般步骤：
1. 设计框架
2. 设计骨干网络
Unet网络设计的步骤：
1. 设计Unet网络工厂模式
2. 设计编解码结构
3. 设计卷积模块
4. unet实例模块
Unet网络最重要的特征：
1. 编解码结构。
2. 解码结构，比FCN更加完善，采用连接方式。
3. 本质是一个框架，编码部分可以使用很多图像分类网络。
示例代码：
import torch
import torch.nn as nn
class Unet(nn.Module):
#初始化参数：Encoder,Decoder,bridge
#bridge默认值为无，如果有参数传入，则用该参数替换None
def __init__(self,Encoder,Decoder,bridge = None):
super(Unet,self).__init__()
self.encoder = Encoder(encoder_blocks)
self.decoder = Decoder(decoder_blocks)
self.bridge = bridge
def forward(self,x):
res = self.encoder(x)
out,skip = res[0],res[1,:]
if bridge is not None:
out = bridge(out)
out = self.decoder(out,skip)
return out
#设计编码模块
class Encoder(nn.Module):
def __init__(self,blocks):
super(Encoder,self).__init__()
#assert:断言函数，避免出现参数错误
assert len(blocks) > 0
#nn.Modulelist()：模型列表，所有的参数可以纳入网络，但是没有forward函数
self.blocks = nn.Modulelist(blocks)
def forward(self,x):
skip = []
for i in range(len(self.blocks) - 1):
x = self.blocks[i](x)
skip.append(x)
res = [self.block[i+1](x)]
#列表之间可以通过+号拼接
res += skip
return res
#设计Decoder模块
class Decoder(nn.Module):
def __init__(self,blocks):
super(Decoder, self).__init__()
assert len(blocks) > 0
self.blocks = nn.Modulelist(blocks)
def ceter_crop(self,skips,x):
_,_,height1,width1 = skips.shape()
_,_,height2,width2 = x.shape()
#对图像进行剪切处理，拼接的时候保持对应size参数一致
ht,wt = min(height1,height2),min(width1,width2)
dh1 = (height1 - height2)//2 if height1 > height2 else 0
dw1 = (width1 - width2)//2 if width1 > width2 else 0
dh2 = (height2 - height1)//2 if height2 > height1 else 0
dw2 = (width2 - width1)//2 if width2 > width1 else 0
return skips[:,:,dh1:(dh1 + ht),dw1:(dw1 + wt)],\
x[:,:,dh2:(dh2 + ht),dw2 : (dw2 + wt)]
def forward(self, skips,x,reverse_skips = True):
assert len(skips) == len(blocks) - 1
if reverse_skips is True:
skips = skips[: : -1]
x = self.blocks[0](x)
for i in range(1, len(self.blocks)):
skip = skips[i-1]
x = torch.cat(skip,x,1)
x = self.blocks[i](x)
return x
#定义了一个卷积block
def unet_convs(in_channels,out_channels,padding = 0):
#nn.Sequential：与Modulelist相比，包含了forward函数
return nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernal_size = 3, padding = padding, bias = False),
nn.BatchNorm2d(outchannels),
nn.ReLU(inplace = True),
nn.Conv2d(in_channels, out_channels, kernal_size=3, padding=padding, bias=False),
nn.BatchNorm2d(outchannels),
nn.ReLU(inplace=True),
)
#实例化Unet模型
def unet(in_channels,out_channels):
encoder_blocks = [unet_convs(in_channels, 64),\
nn.Sequential(nn.Maxpool2d(kernal_size = 2, stride = 2, ceil_mode = True),\
unet_convs(64,128)), \
nn.Sequential(nn.Maxpool2d(kernal_size=2, stride=2, ceil_mode=True), \
unet_convs(128, 256)),
nn.Sequential(nn.Maxpool2d(kernal_size=2, stride=2, ceil_mode=True), \
unet_convs(256, 512)),
]
bridge = nn.Sequential(unet_convs(512, 1024))
decoder_blocks = [nn.conTranpose2d(1024, 512), \
nn.Sequential(unet_convs(1024, 512),
nn.conTranpose2d(512, 256)),\
nn.Sequential(unet_convs(512, 256),
nn.conTranpose2d(256, 128)), \
nn.Sequential(unet_convs(512, 256),
nn.conTranpose2d(256, 128)), \
nn.Sequential(unet_convs(256, 128),
nn.conTranpose2d(128, 64))
]
return Unet(encoder_blocks,decoder_blocks,bridge)
补充知识：Pytorch搭建U-Net网络
U-Net: Convolutional Networks for Biomedical Image Segmentation
import torch.nn as nn
import torch
from torch import autograd
from torchsummary import summary
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=0),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
nn.Conv2d(out_ch, out_ch, 3, padding=0),
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.ConvTranspose2d(512, 256, 2, stride=2)
self.conv7 = DoubleConv(512, 256)
self.up8 = nn.ConvTranspose2d(256, 128, 2, stride=2)
self.conv8 = DoubleConv(256, 128)
self.up9 = nn.ConvTranspose2d(128, 64, 2, stride=2)
self.conv9 = DoubleConv(128, 64)
self.conv10 = nn.Conv2d(64, out_ch, 1)
def forward(self, x):
c1 = self.conv1(x)
crop1 = c1[:,:,88:480,88:480]
p1 = self.pool1(c1)
c2 = self.conv2(p1)
crop2 = c2[:,:,40:240,40:240]
p2 = self.pool2(c2)
c3 = self.conv3(p2)
crop3 = c3[:,:,16:120,16:120]
p3 = self.pool3(c3)
c4 = self.conv4(p3)
crop4 = c4[:,:,4:60,4:60]
p4 = self.pool4(c4)
c5 = self.conv5(p4)
up_6 = self.up6(c5)
merge6 = torch.cat([up_6, crop4], dim=1)
c6 = self.conv6(merge6)
up_7 = self.up7(c6)
merge7 = torch.cat([up_7, crop3], dim=1)
c7 = self.conv7(merge7)
up_8 = self.up8(c7)
merge8 = torch.cat([up_8, crop2], dim=1)
c8 = self.conv8(merge8)
up_9 = self.up9(c8)
merge9 = torch.cat([up_9, crop1], dim=1)
c9 = self.conv9(merge9)
c10 = self.conv10(c9)
out = nn.Sigmoid()(c10)
return out
if __name__=="__main__":
test_input=torch.rand(1, 1, 572, 572)
model=Unet(in_ch=1, out_ch=2)
summary(model, (1,572,572))
ouput=model(test_input)
print(ouput.size())
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 570, 570] 640
BatchNorm2d-2 [-1, 64, 570, 570] 128
ReLU-3 [-1, 64, 570, 570] 0
Conv2d-4 [-1, 64, 568, 568] 36,928
BatchNorm2d-5 [-1, 64, 568, 568] 128
ReLU-6 [-1, 64, 568, 568] 0
DoubleConv-7 [-1, 64, 568, 568] 0
MaxPool2d-8 [-1, 64, 284, 284] 0
Conv2d-9 [-1, 128, 282, 282] 73,856
BatchNorm2d-10 [-1, 128, 282, 282] 256
ReLU-11 [-1, 128, 282, 282] 0
Conv2d-12 [-1, 128, 280, 280] 147,584
BatchNorm2d-13 [-1, 128, 280, 280] 256
ReLU-14 [-1, 128, 280, 280] 0
DoubleConv-15 [-1, 128, 280, 280] 0
MaxPool2d-16 [-1, 128, 140, 140] 0
Conv2d-17 [-1, 256, 138, 138] 295,168
BatchNorm2d-18 [-1, 256, 138, 138] 512
ReLU-19 [-1, 256, 138, 138] 0
Conv2d-20 [-1, 256, 136, 136] 590,080
BatchNorm2d-21 [-1, 256, 136, 136] 512
ReLU-22 [-1, 256, 136, 136] 0
DoubleConv-23 [-1, 256, 136, 136] 0
MaxPool2d-24 [-1, 256, 68, 68] 0
Conv2d-25 [-1, 512, 66, 66] 1,180,160
BatchNorm2d-26 [-1, 512, 66, 66] 1,024
ReLU-27 [-1, 512, 66, 66] 0
Conv2d-28 [-1, 512, 64, 64] 2,359,808
BatchNorm2d-29 [-1, 512, 64, 64] 1,024
ReLU-30 [-1, 512, 64, 64] 0
DoubleConv-31 [-1, 512, 64, 64] 0
MaxPool2d-32 [-1, 512, 32, 32] 0
Conv2d-33 [-1, 1024, 30, 30] 4,719,616
BatchNorm2d-34 [-1, 1024, 30, 30] 2,048
ReLU-35 [-1, 1024, 30, 30] 0
Conv2d-36 [-1, 1024, 28, 28] 9,438,208
BatchNorm2d-37 [-1, 1024, 28, 28] 2,048
ReLU-38 [-1, 1024, 28, 28] 0
DoubleConv-39 [-1, 1024, 28, 28] 0
ConvTranspose2d-40 [-1, 512, 56, 56] 2,097,664
Conv2d-41 [-1, 512, 54, 54] 4,719,104
BatchNorm2d-42 [-1, 512, 54, 54] 1,024
ReLU-43 [-1, 512, 54, 54] 0
Conv2d-44 [-1, 512, 52, 52] 2,359,808
BatchNorm2d-45 [-1, 512, 52, 52] 1,024
ReLU-46 [-1, 512, 52, 52] 0
DoubleConv-47 [-1, 512, 52, 52] 0
ConvTranspose2d-48 [-1, 256, 104, 104] 524,544
Conv2d-49 [-1, 256, 102, 102] 1,179,904
BatchNorm2d-50 [-1, 256, 102, 102] 512
ReLU-51 [-1, 256, 102, 102] 0
Conv2d-52 [-1, 256, 100, 100] 590,080
BatchNorm2d-53 [-1, 256, 100, 100] 512
ReLU-54 [-1, 256, 100, 100] 0
DoubleConv-55 [-1, 256, 100, 100] 0
ConvTranspose2d-56 [-1, 128, 200, 200] 131,200
Conv2d-57 [-1, 128, 198, 198] 295,040
BatchNorm2d-58 [-1, 128, 198, 198] 256
ReLU-59 [-1, 128, 198, 198] 0
Conv2d-60 [-1, 128, 196, 196] 147,584
BatchNorm2d-61 [-1, 128, 196, 196] 256
ReLU-62 [-1, 128, 196, 196] 0
DoubleConv-63 [-1, 128, 196, 196] 0
ConvTranspose2d-64 [-1, 64, 392, 392] 32,832
Conv2d-65 [-1, 64, 390, 390] 73,792
BatchNorm2d-66 [-1, 64, 390, 390] 128
ReLU-67 [-1, 64, 390, 390] 0
Conv2d-68 [-1, 64, 388, 388] 36,928
BatchNorm2d-69 [-1, 64, 388, 388] 128
ReLU-70 [-1, 64, 388, 388] 0
DoubleConv-71 [-1, 64, 388, 388] 0
Conv2d-72 [-1, 2, 388, 388] 130
================================================================
Total params: 31,042,434
Trainable params: 31,042,434
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 1.25
Forward/backward pass size (MB): 3280.59
Params size (MB): 118.42
Estimated Total Size (MB): 3400.26
----------------------------------------------------------------
torch.Size([1, 2, 388, 388])
以上这篇使用pytorch实现论文中的unet网络就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持脚本之家。