使用pytorch搭建Unet
import torch
import torch.nn as nn
from torch.nn import functional as F
# 上采样方式使用 nn.UpsamplingBilinear2d(scale_factor=2)
def conv2_relu(in_chan, out_chan, pd):
return nn.Sequential(
nn.Conv2d(in_channels=in_chan, out_channels=out_chan, kernel_size=3, padding=pd),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=out_chan, out_channels=out_chan, kernel_size=3, padding=pd),
nn.ReLU(inplace=True)
)
def conv3_relu(in_chan, out_chan, pd):
return nn.Sequential(
nn.Conv2d(in_channels=in_chan, out_channels=out_chan, kernel_size=3, padding=pd),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=out_chan, out_channels=out_chan, kernel_size=3, padding=pd),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels=out_chan, out_channels=out_chan, kernel_size=3, padding=pd),
nn.ReLU(inplace=True)
)
class U_Net(nn.Module):
def __init__(self, num_cls):
super(U_Net, self).__init__()
self.num_cls = num_cls
# 编码器卷积
self.conv_1 = conv2_relu(in_chan=3, out_chan=64, pd=1)
self.conv_2 = conv2_relu(in_chan=64, out_chan=128, pd=1)
self.conv_3 = conv3_relu(in_chan=128, out_chan=256, pd=1)
self.conv_4 = conv3_relu(in_chan=256, out_chan=512, pd=1)
self.conv_5 = conv3_relu(in_chan=512, out_chan=512, pd=1)
# 解码器卷积
self.conv_6 = conv2_relu(in_chan=1024, out_chan=512, pd=1)
self.conv_7 = conv2_relu(in_chan=768, out_chan=256, pd=1)
self.conv_8 = conv2_relu(in_chan=384, out_chan=128, pd=1)
self.conv_9 = conv2_relu(in_chan=192, out_chan=64, pd=1)
self.conv_final = nn.Conv2d(in_channels=64, out_channels=self.num_cls, kernel_size=1)
self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
self.upsample = nn.UpsamplingBilinear2d(scale_factor=2)
def forward(self,x):
# 编码器部分 下采样
# x:(N, 3, 416, 416)
f1 = self.conv_1(x) # (N, 64, 416, 416)
x = self.maxpool(f1) # (N, 64, 208, 208)
f2 = self.conv_2(x) # (N, 128, 208, 208)
x = self.maxpool(f2) # (N, 128, 104, 104)
f3 = self.conv_3(x) # (N, 256, 104, 104)
x = self.maxpool(f3) # (N, 256, 52, 52)
f4 = self.conv_4(x) # (N, 512, 52, 52)
x = self.maxpool(f4) # (N, 512, 26, 26)
f5 = self.conv_5(x) # (N, 512, 26, 26)
# 解码器部分 上采样
p5_up = self.upsample(f5) # (N, 512, 52, 52)
p4 = torch.cat([f4, p5_up], dim=1) # (N, 1024, 52, 52); 512+512=1024
p4 = self.conv_6(p4) # (N, 512, 52, 52)
p4_up = self.upsample(p4) # (N, 512, 104, 104)
p3 = torch.cat([f3, p4_up], dim=1) # (N, 768, 104, 104); 256+512=768
p3 = self.conv_7(p3) # (N, 256, 104, 104)
p3_up = self.upsample(p3) # (N, 256, 208, 208)
p2 = torch.cat([f2, p3_up], dim=1) # (N, 384, 208, 208); 128+256=384
p2 = self.conv_8(p2) # (N, 128, 208, 208)
p2_up = self.upsample(p2) # (N, 128, 416, 416)
p1 = torch.cat([f1, p2_up], dim=1) # (N, 192, 416, 416)
p1 = self.conv_9(p1) # (N, 64, 416, 416)
p1 = self.conv_final(p1) # (N, cls_num, 416, 416)
return p1
model = U_Net(num_cls=9)
data = torch.rand([1, 3, 416, 416]).float()
out = model(data)
print('out.shape:', out.shape)
输出结果:
out.shape: torch.Size([1, 9, 416, 416])