import torch
from torch import nn
from torch.nn import functional as F
class Conv_Block(nn.Module):
def __init__(self,in_channel,out_channel):
super(Conv_Block, self).__init__()
self.layer=nn.Sequential(
nn.Conv2d(in_channel,out_channel,3,1,1,padding_mode='reflect',bias=False),
nn.BatchNorm2d(out_channel),
nn.Dropout2d(0.3),
nn.LeakyReLU(),
nn.Conv2d(out_channel, out_channel, 3, 1, 1, padding_mode='reflect', bias=False),
nn.BatchNorm2d(out_channel),
nn.Dropout2d(0.3),
nn.LeakyReLU()
)
def forward(self,x):
return self.layer(x)
class DownSample(nn.Module):
def __init__(self,channel):
super(DownSample, self).__init__()
self.layer=nn.Sequential(
nn.MaxPool2d(2),
nn.Conv2d(channel,channel,3,1,1,padding_mode='reflect',bias=False),
nn.BatchNorm2d(channel),
nn.LeakyReLU()
)
def forward(self,x):
return self.layer(x)
class UpSample(nn.Module):
def __init__(self,channel):
super(UpSample, self).__init__()
self.layer=nn.Conv2d(channel,channel//2,1,1)
def forward(self,x,feature_map):
up=F.interpolate(x,scale_factor=2,mode='nearest')#也可以使bilinear
out=self.layer(up)
return torch.cat((out,feature_map),dim=1)
class UNet(nn.Module):
def __init__(self,num_classes):
super(UNet, self).__init__()
self.c1=Conv_Block(3,64)
self.d1=DownSample(64)
self.c2=Conv_Block(64,128)
self.d2=DownSample(128)
self.c3=Conv_Block(128,256)
self.d3=DownSample(256)
self.c4=Conv_Block(256,512)
self.d4=DownSample(512)
self.c5=Conv_Block(512,1024)
self.u1=UpSample(1024)
self.c6=Conv_Block(1024,512)
self.u2 = UpSample(512)
self.c7 = Conv_Block(512, 256)
self.u3 = UpSample(256)
self.c8 = Conv_Block(256, 128)
self.u4 = UpSample(128)
self.c9 = Conv_Block(128, 64)
self.out=nn.Conv2d(64,num_classes,3,1,1)
def forward(self,x):
R1=self.c1(x)
R2=self.c2(self.d1(R1))
R3 = self.c3(self.d2(R2))
R4 = self.c4(self.d3(R3))
R5 = self.c5(self.d4(R4))
O1=self.c6(self.u1(R5,R4))
O2 = self.c7(self.u2(O1, R3))
O3 = self.c8(self.u3(O2, R2))
O4 = self.c9(self.u4(O3, R1))
return self.out(O4)
if __name__ == '__main__':
x=torch.randn(2,3,256,256)
net=UNet()
print(net(x).shape)
DownSample模块是一个下采样模块(max pool),UpSample模块是一个上采样模块(up conv)。
其中UpSample模块使用双线性插值法的结果
UpSample模块使用邻近0插值法的结果
DownSample包括max pool+2次的conv,
UpSample模块 包括
up=F.interpolate(x,scale_factor=2,mode='nearest')就是上采样插值扩张
然后就是conv卷积
Conv_Block层次 卷积层(2次的conv)
- nn.Conv2d(in_channel,out_channel,3,1,1,padding_mode=‘reflect’,bias=False):这是一个卷积层,其中in_channel表示输入图像的通道数,out_channel表示输出图像的通道数,3表示卷积核的大小为3*3,1表示步长为1,padding_mode='reflect’表示使用反射填充,即在边缘处进行填充,bias=False表示不使用偏置。
- nn.BatchNorm2d(out_channel):这是一个批标准化层,用于规范化每个输出通道的均值和方差,加快网络收敛速度。
- nn.Dropout2d(0.3):这是一个dropout层,用于防止过拟合,其中0.3表示随机失活的比例。
- nn.LeakyReLU():这是一个激活函数,用于提高网络的非线性拟合能力,其具体函数为f(x)=max(ax,x),其中a为负斜率系数。
欧克,