unet代码以及注释，免费使用

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( # 3: 卷积核大小为 3x3。1: 步长为 1。1: padding填充大小为 1。padding_mode='reflect': 使用反射填充模式，即通过对输入进行反射填充来处理边界。 nn.Conv2d(in_channel,out_channel,3,1,1,padding_mode='reflect',bias=False), nn.BatchNorm2d(out_channel), # 添加了一个丢弃层，以概率 0.3 随机将输入张量中的部分元素设置为零，有助于防止过拟合 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.Conv2d(channel,channel,3,2,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') out=self.layer(up) return torch.cat((out,feature_map),dim=1) class UNet(nn.Module): def __init__(self): 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,3,3,1,1) self.TH = nn.Sigmoid() 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.TH(self.out(o4)) return self.out(o4) if __name__ == '__main__': x=torch.randn(2,3,256,256) net=UNet() print(net(x).shape)