图像分割之U-Net、U2-Net及其Pytorch代码构建(图像分割 unet)

图像分割之U-Net、U2-Net及其Pytorch代码构建 图像分割之U-Net、U2 -Net及其Pytorch代码构建1、图像分割

推荐整理分享图像分割之U-Net、U2-Net及其Pytorch代码构建(图像分割 unet)，希望有所帮助，仅作参考，欢迎阅读内容。

文章相关热门搜索词:图像分割的经典算法fcn,图像分割的经典算法fcn,图像分割ncut,图像分割csdn,图像分割otsu,图像分割otsu,图像分割otsu,图像分割 unet,内容如对您有帮助，希望把文章链接给更多的朋友！

图像分割就是把图像分成若干个特定的、具有独特性质的区域并提出感兴趣目标的技术和过程。

做法便是对图片中的每一个像素进行分类。

在自动驾驶、自动抠图、医疗影像等领域有着比较广泛的应用。

图像分割大致可分为以下三类：

普通分割：将不同分属不同物体的像素区域分开。比如前景和背景分割开，狗的区域和猫的区域与背景分割开。语义分割：在普通分割的基础上，分类出每一块区域的语义（即这块区域是什么物体）。如把画面中的所有物体都指出他们各自的类别。实例分割：在语义分割的基础上，给每一个物体编号。如这个是该画面中的狗A，那个是画面中的狗B。普通分割语义分割实例分割

可以看出，图像分割是由一张图片到另一张图片。因此，神经网络的输入是图片，输出也是同样的图片，Encoder-Decoder的结构是合适的。U-Net、U2 -Net可作为语义分割使用，可以按照生成图像的方式，生成分割图。也可以按通道划分类，每一个通道就是一个类别，使用sigmoid激活。

2、U-Net

U-Net即使用Encoder-Decoder的结构，首先下采样，然后上采样，中间每一级由残差组成。

则可构建网络的代码如下：

首先是卷积层，可以看出，网络在每一级，均有两层卷积组成。因此构建卷积层如下：

from torch import nnimport torchclass ConvolutionLayer(nn.Module): def __init__(self, in_channels, out_channels): """ 卷积层 :param in_channels: 输入通道 :param out_channels: 输出通道 """ super(ConvolutionLayer, self).__init__() self.layer = nn.Sequential( # 卷积层 nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), stride=(1, 1), padding=1, bias=False), nn.BatchNorm2d(out_channels), # BN层 nn.ReLU(), # 激活 nn.Conv2d(out_channels, out_channels, kernel_size=(3, 3), stride=(1, 1), padding=1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU(), ) def forward(self, x): return self.layer(x)

同时与图示不同的地方在于，使用了Padding，以免图片在卷积中的尺寸缩小。这样，横向的灰色箭头可以直接使用cat进行两个特征图的拼接。

模型图中，红色箭头的max pool 2×2，使用的是池化窗口为2×2的最大值池化。这里的目的是进行下采样，因此可以定义一个下采样如下：

class DownSample(nn.Module): def __init__(self,): """ 最大池化层构成的下采样，池化窗口为2×2 """ super(DownSample, self).__init__() self.layer = nn.MaxPool2d(kernel_size=2, stride=2) def forward(self, x): return self.layer(x)

模型图中，绿色箭头的up-conv 2×2，使用的是反卷积。这里的目的是进行上采样，因此可以定义一个上采样如下：

class UpSample(nn.Module): def __init__(self, in_channels): """ 反卷积，上采样，通道数将会减半， :param in_channels: 输入通道数 """ super(UpSample, self).__init__() self.layer = nn.Sequential( nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=(2, 2), stride=(2, 2)), nn.LeakyReLU(), ) def forward(self, x): return self.layer(x)

首先定义各个网络层：

class UNet(nn.Module): def __init__(self, in_channels, out_channels): super(UNet, self).__init__() self.conv1 = ConvolutionLayer(in_channels, 64) # 三通道拓展至64通道 self.down1 = DownSample() # 下采样至1/2 self.conv2 = ConvolutionLayer(64, 128) # 64通道==>128通道 self.down2 = DownSample() # 下采样至1/4 self.conv3 = ConvolutionLayer(128, 256) # 128通道==>256通道 self.down3 = DownSample() # 下采样至1/8 self.conv4 = ConvolutionLayer(256, 512) # 256通道==>512通道 self.down4 = DownSample() # 下采样至1/16 self.conv5 = ConvolutionLayer(512, 1024) # 512通道==>1024通道 self.up1 = UpSample(1024) # 上采样至1/8 self.conv6 = ConvolutionLayer(1024, 512) # 1024通道==>512通道 self.up2 = UpSample(512) # 上采样至1/4 self.conv7 = ConvolutionLayer(512, 256) # 512通道==>256通道 self.up3 = UpSample(256) # 上采样至1/2 self.conv8 = ConvolutionLayer(256, 128) # 256通道==>128通道 self.up4 = UpSample(128) # 上采样至1/1 self.conv9 = ConvolutionLayer(128, 64) # 128通道==>64通道 self.predict = nn.Sequential( # 输出层，由sigmoid函数激活 nn.Conv2d(64, out_channels, kernel_size=(3,3), stride=(1,1), padding=1), nn.Sigmoid() ) def forward(self, image_tensor):pass

对应于模型图如下：

class UNet(nn.Module): def __init__(self, in_channels, out_channels): super(UNet, self).__init__() """ ...... """ def forward(self, x): """下采样""" x1 = self.conv1(x) # ===> 1/1 64 d1 = self.down1(x1) # ===> 1/2 64 x2 = self.conv2(d1) # ===> 1/2 128 d2 = self.down2(x2) # ===> 1/4 128 x3 = self.conv3(d2) # ===> 1/4 256 d3 = self.down3(x3) # ===> 1/8 256 x4 = self.conv4(d3) # ===> 1/8 512 d4 = self.down4(x4) # ===> 1/16 512 x5 = self.conv5(d4) # ===> 1/16 1024 """上采样""" up1 = self.up1(x5) # ===> 1/8 512 x6 = self.conv6(torch.cat((x4, up1), dim=1)) # ===> 1/8 512 up2 = self.up2(x6 # ===> 1/4 256 x7 = self.conv7(torch.cat((x3, up2), dim=1)) # ===> 1/4 256 up3 = self.up3(x7) # ===> 1/2 128 x8 = self.conv8(torch.cat((x2, up3), dim=1)) # ===> 1/2 128 up4 = self.up4(x8) # ===> 1/1 64 x9 = self.conv9(torch.cat((x1, up4), dim=1)) # ===> 1/1 64 mask = self.predict(x9) # ===> 1/1 out_channels return mask

以一张512×512的3通道图片为例，其张量的形状为（1,3,512,512），经过conv1得到x1 （1, 64, 512, 512），下采样至（1, 64, 256, 256）；经过conv2得到x2 （1, 128, 256, 256），下采样至（1, 128, 128, 128）；经过conv3得到x3 （1, 256, 128, 128），下采样至（1, 256, 64, 64）；经过conv4得到x4 （1, 512, 64, 64），下采样至（1, 512, 32, 32）；经过conv5得到x5 （1, 1024, 32, 32）。下采样过程完成，开始上采样还原至原始图片大小。

x5经过up1得到up1 （1, 512, 64, 64）,同x4 拼接（cat）在一起 组成（1, 1024, 64, 64）的张量，经过conv6得到x6（1, 512, 64, 64）;

x6经过up2得到up2 （1, 256, 128, 128）,同x3 拼接在一起 组成（1, 512, 128, 128）的张量，经过conv7得到x7（1, 256, 128, 128）;

x7经过up3得到up3 （1, 128, 256, 256）,同x2 拼接在一起 组成（1, 256, 256, 256）的张量，经过conv8得到x8（1, 128, 256, 256）;

x8经过up4得到up4 （1, 64, 512, 512）,同x1 拼接在一起 组成（1, 128, 512, 512）的张量，经过conv6得到x9（1, 64, 512, 512）;

最后，x9经过预测层predict输出，得到分割图mask。

以drive数据集为例训练网络，数据示例如下。

标签如下：

输入数据为3通道的图片，而输出数据为1通道的二值图。一张图片的原始尺寸是565×584

可以在原始图像中随机裁剪256×256大小的图片，进行训练，而在使用时，图像尺寸只要是16的倍数即可。

定义数据加载函数如下：

import torchimport randomimport cv2from torch.utils.data import Datasetclass DriveDataset(Dataset): def __init__(self,root='data/training'): super(DriveDataset, self).__init__() self.dataset = [] start = 20 for i in range(1, 21): # 按照一一对应的原则，加载图像和标签的路径 image_path = f'{root}/images/{i+start}_training.tif' label_path = f'{root}/1st_manual/{i + start}_manual1.gif' self.dataset.append((image_path, label_path)) def __len__(self): return len(self.dataset) def __getitem__(self, item): image_path, label_path = self.dataset[item] # 获取图像路径 image = cv2.imread(image_path) # 图片 video = cv2.VideoCapture(label_path) _, mask_label = video.read() # 读取标签掩码图 image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) mask_label = cv2.cvtColor(mask_label, cv2.COLOR_BGR2GRAY) # 转换至单通道图 """随即裁剪256×256的图幅，图片和标签裁剪相同的位置""" h, w = mask_label.shape w = random.randint(0, w-256) h = random.randint(0, h-256) image = image[h:h+256, w:w+256] mask_label = mask_label[h:h + 256, w:w + 256] """转换至tensor""" image = torch.from_numpy(image).float().permute(2, 0, 1)/255 mask_label = torch.from_numpy(mask_label).unsqueeze(0).float()/255 return image, mask_label

读取相对应的图片和标签，转换为张量，供网络学习。其中，标签的读取使用了OpenCV的视频捕获（VideoCapture）读取首帧完成标签的数据加载。

定义训练器如下：

from torch import nnimport torchfrom torch.utils.data import DataLoaderfrom torchvision.utils import save_imagefrom u_net import UNetfrom dataset import DriveDatasetimport osclass Trainer: def __init__(self): self.device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu") # 设置设备 self.net = UNet(3, 1).to(self.device) # 实例U-Net if os.path.exists('unet.pth'): # 加载权重，如果存在的话 self.net.load_state_dict(torch.load('unet.pth', map_location='cpu')) self.dataset = DriveDataset() # 实例数据集 self.data_loader = DataLoader(self.dataset, 3, True, drop_last=True) # 实例数据加载器 self.loss_func = nn.BCELoss() # 实例二值交叉熵 self.optimizer = torch.optim.Adam(self.net.parameters()) # 实例adam优化器 def train(self): # 训练 for epoch in range(100000): # 迭代epoch for i, (image, target) in enumerate(self.data_loader): image = image.to(self.device) target = target.to(self.device) out = self.net(image) # 预测 loss = self.loss_func(out, target) # 计算损失 self.optimizer.zero_grad() # 清空梯度 loss.backward() # 反向传播 self.optimizer.step() # 优化 print(epoch, loss.item()) if epoch % 5 == 0: torch.save(self.net.state_dict(),'unet.pth') save_image([image[0], target[0].expand(3, 256, 256), out[0].expand(3, 256, 256)], f'{epoch}.jpg',normalize=True,range=(0,1))

二值交叉熵做损失，adam优化器优化网络。

class Trainer: """ ...... """if __name__ == '__main__': trainer = Trainer() trainer.train()

训练过程见下图。左边为原图，中间为标签，右边为网络预测值

epochimages012

完整代码：https://github.com/HibikiJie/UNetAndU2Net

3、U2-Net

而U2-Net，就是U-Net的堆叠，类似于，将U-Net中的conv块，替换成完整的U-Net网络。

其网络图如下：

其中EN_1与De_1一致，EN_2与De_2一致，EN_3与De_3一致，EN_4与De_4一致，EN_5、En6和De_5一致。

先分别定义，EN_1、EN_2、EN_3、EN_4、EN_5为UNet1、UNet2、UNet3、UNet4、UNet5.

首先定义UNet1：

注意到，图中的白色的方块示意的，卷积使用到了dilation参数，因此，定义ConvolutionLayer为：

import torchimport torch.nn as nnimport torch.nn.functional as Fclass ConvolutionLayer(nn.Module): def __init__(self, in_channels, out_channels, dilation=1): super(ConvolutionLayer, self).__init__() self.layer = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), padding=1 * dilation, dilation=(1 * dilation, 1 * dilation)), # 卷积 nn.BatchNorm2d(out_channels), # BN nn.ReLU(inplace=True) # 激活函数 ) def forward(self, x): return self.layer(x)

卷积层由Conv、BN、ReLU构成。

上采样使用机器学习算法，由双线性插值法完成上采样：

def upsample_like(src, tar): src = F.upsample(src, size=tar.shape[2:], mode='bilinear') return src

该方法，将使src上采样至tar相同的尺寸大小。

而下采样同样使用最大池化完成，这里可以使用与U-Net相同的代码。

因此，UNet1：

class UNet1(nn.Module): def __init__(self, in_channels, mid_channels, out_channels): super(UNet1, self).__init__() self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1) self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1) self.down1 = DownSample() self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down2 = DownSample() self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down3 = DownSample() self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down4 = DownSample() self.conv5 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down5 = DownSample() self.conv6 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.conv7 = ConvolutionLayer(mid_channels, mid_channels, dilation=2) self.conv8 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv9 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv10 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv11 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv12 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv13 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1) def forward(self, x): """下采样，编码encode的过程""" x0 = self.conv0(x) x1 = self.conv1(x0) d1 = self.down1(x1) x2 = self.conv2(d1) d2 = self.down2(x2) x3 = self.conv3(d2) d3 = self.down3(x3) x4 = self.conv4(d3) d4 = self.down4(x4) x5 = self.conv5(d4) d5 = self.down5(x5) x6 = self.conv6(d5) x7 = self.conv7(x6)"""上采样，解码decode的过程""" x8 = self.conv8(torch.cat((x7, x6), dim=1)) up1 = upsample_like(x8, x5) x9 = self.conv9(torch.cat((up1, x5), dim=1)) up2 = upsample_like(x9, x4) x10 = self.conv10(torch.cat((up2, x4), dim=1)) up3 = upsample_like(x10, x3) x11 = self.conv11(torch.cat((up3, x3), dim=1)) up4 = upsample_like(x11, x2) x12 = self.conv12(torch.cat((up4, x2), dim=1)) up5 = upsample_like(x12, x1) x13 = self.conv13(torch.cat((up5, x1), dim=1)) return x13 + x0

按照上图所示的方式编码，可见，与写UNet的代码是非常类似的。可以对比着看。可见，U2-Net是U-Net的堆叠。

于是类似的，UNet2的代码为:

class UNet2(nn.Module): def __init__(self, in_channels, mid_channels, out_channels): super(UNet2, self).__init__() self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1) self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1) self.down1 = DownSample() self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down2 = DownSample() self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down3 = DownSample() self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down4 = DownSample() self.conv5 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.conv6 = ConvolutionLayer(mid_channels, mid_channels, dilation=2) self.conv7 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv8 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv9 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv10 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv11 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1) def forward(self, x): """encode""" x0 = self.conv0(x) x1 = self.conv1(x0) d1 = self.down1(x1) x2 = self.conv2(d1) d2 = self.down2(x2) x3 = self.conv3(d2) d3 = self.down3(x3) x4 = self.conv4(d3) d4 = self.down4(x4) x5 = self.conv5(d4) x6 = self.conv6(x5)"""decode""" x7 = self.conv7(torch.cat((x6, x5), dim=1)) up1 = upsample_like(x7, x4) x8 = self.conv8(torch.cat((up1, x4), dim=1)) up2 = upsample_like(x8, x3) x9 = self.conv9(torch.cat((up2, x3), dim=1)) up3 = upsample_like(x9, x2) x10 = self.conv10(torch.cat((up3, x2), dim=1)) up4 = upsample_like(x10, x1) x11 = self.conv11(torch.cat((up4, x1), dim=1)) return x11 + x0

UNet3为：

class UNet3(nn.Module): def __init__(self, in_channels, mid_channels, out_channels): super(UNet3, self).__init__() self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1) self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1) self.down1 = DownSample() self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down2 = DownSample() self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down3 = DownSample() self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.conv5 = ConvolutionLayer(mid_channels, mid_channels, dilation=2) self.conv6 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv7 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv8 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv9 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1) def forward(self, x): """encode""" x0 = self.conv0(x) x1 = self.conv1(x0) d1 = self.down1(x1) x2 = self.conv2(d1) d2 = self.down2(x2) x3 = self.conv3(d2) d3 = self.down3(x3) x4 = self.conv4(d3) x5 = self.conv5(x4)"""decode""" x6 = self.conv6(torch.cat((x5, x4), dim=1)) up1 = upsample_like(x6, x3) x7 = self.conv7(torch.cat((up1, x3), dim=1)) up2 = upsample_like(x7, x2) x8 = self.conv8(torch.cat((up2, x2), dim=1)) up3 = upsample_like(x8, x1) x9 = self.conv9(torch.cat((up3, x1), dim=1)) return x9 + x0

UNet4为:

class UNet4(nn.Module): def __init__(self, in_channels, mid_channels, out_channels): super(UNet4, self).__init__() self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1) self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1) self.down1 = DownSample() self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down2 = DownSample() self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=2) self.conv5 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv6 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv7 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1) def forward(self, x): """encode""" x0 = self.conv0(x) x1 = self.conv1(x0) d1 = self.down1(x1) x2 = self.conv2(d1) d2 = self.down2(x2) x3 = self.conv3(d2) x4 = self.conv4(x3) """decode""" x5 = self.conv5(torch.cat((x4, x3), dim=1)) up1 = upsample_like(x5, x2) x6 = self.conv6(torch.cat((up1, x2), dim=1)) up2 = upsample_like(x6, x1) x7 = self.conv7(torch.cat((up2, x1), dim=1)) return x7 + x0

UNet5为:

class UNet5(nn.Module): def __init__(self, in_channels, mid_channels, out_channels): super(UNet5, self).__init__() self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1) self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1) self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=2) self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=4) self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=8) self.conv5 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=4) self.conv6 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=2) self.conv7 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1) def forward(self, x): x0 = self.conv0(x) x1 = self.conv1(x0) x2 = self.conv2(x1) x3 = self.conv3(x2) x4 = self.conv4(x3) x5 = self.conv5(torch.cat((x4, x3), dim=1)) x6 = self.conv6(torch.cat((x5, x2), dim=1)) x7 = self.conv7(torch.cat((x6, x1), dim=1)) return x7 + x0

于是将

UNet1、UNet2、UNet3、UNet4、UNet5.组装成为U2-Net

再看一下网络结构图：

其中EN_1与De_1一致，EN_2与De_2一致，EN_3与De_3一致，EN_4与De_4一致，EN_5、En6和De_5一致。

先分别定义，EN_1、EN_2、EN_3、EN_4、EN_5为UNet1、UNet2、UNet3、UNet4、UNet5.

于是EN_1与De_1使用UNet1；

EN_2与De_2使用UNet2；

EN_3与De_3使用UNet3；

EN_4与De_4使用UNet4；

EN_5、EN_6、De_5使用UNet1。

故，构建网络U2-Net：

class U2Net(nn.Module): def __init__(self, in_channels=3, out_channels=1): super(U2Net, self).__init__() self.en_1 = UNet1(in_channels, 32, 64) self.down1 = DownSample() self.en_2 = UNet2(64, 32, 128) self.down2 = DownSample() self.en_3 = UNet3(128, 64, 256) self.down3 = DownSample() self.en_4 = UNet4(256, 128, 512) self.down4 = DownSample() self.en_5 = UNet5(512, 256, 512) self.down5 = DownSample() self.en_6 = UNet5(512, 256, 512) # decoder self.de_5 = UNet5(1024, 256, 512) self.de_4 = UNet4(1024, 128, 256) self.de_3 = UNet3(512, 64, 128) self.de_2 = UNet2(256, 32, 64) self.de_1 = UNet1(128, 16, 64) self.side1 = nn.Conv2d(64, out_channels, kernel_size=(3, 3), padding=1) self.side2 = nn.Conv2d(64, out_channels, kernel_size=(3, 3), padding=1) self.side3 = nn.Conv2d(128, out_channels, kernel_size=(3, 3), padding=1) self.side4 = nn.Conv2d(256, out_channels, kernel_size=(3, 3), padding=1) self.side5 = nn.Conv2d(512, out_channels, kernel_size=(3, 3), padding=1) self.side6 = nn.Conv2d(512, out_channels, kernel_size=(3, 3), padding=1) self.out_conv = nn.Conv2d(6, out_channels, kernel_size=(1, 1)) def forward(self, x): # ------encode ------ x1 = self.en_1(x) d1 = self.down1(x1) x2 = self.en_2(d1) d2 = self.down2(x2) x3 = self.en_3(d2) d3 = self.down3(x3) x4 = self.en_4(d3) d4 = self.down4(x4) x5 = self.en_5(d4) d5 = self.down5(x5) x6 = self.en_6(d5) up1 = upsample_like(x6, x5) # ------decode ------ x7 = self.de_5(torch.cat((up1, x5), dim=1)) up2 = upsample_like(x7, x4) x8 = self.de_4(torch.cat((up2, x4), dim=1)) up3 = upsample_like(x8, x3) x9 = self.de_3(torch.cat((up3, x3), dim=1)) up4 = upsample_like(x9, x2) x10 = self.de_2(torch.cat((up4, x2), dim=1)) up5 = upsample_like(x10, x1) x11 = self.de_1(torch.cat((up5, x1), dim=1)) # side output sup1 = self.side1(x11) sup2 = self.side2(x10) sup2 = upsample_like(sup2, sup1) sup3 = self.side3(x9) sup3 = upsample_like(sup3, sup1) sup4 = self.side4(x8) sup4 = upsample_like(sup4, sup1) sup5 = self.side5(x7) sup5 = upsample_like(sup5, sup1) sup6 = self.side6(x6) sup6 = upsample_like(sup6, sup1) sup0 = self.out_conv(torch.cat((sup1, sup2, sup3, sup4, sup5, sup6), dim=1)) return torch.sigmoid(sup0)

U2-Net完整代码：

import torchimport torch.nn as nnimport torch.nn.functional as Fclass ConvolutionLayer(nn.Module): def __init__(self, in_channels, out_channels, dilation=1): super(ConvolutionLayer, self).__init__() self.layer = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), padding=1 * dilation, dilation=(1 * dilation, 1 * dilation)), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True) ) self.conv_s1 = nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), padding=1 * dilation, dilation=(1 * dilation, 1 * dilation)) self.bn_s1 = nn.BatchNorm2d(out_channels) self.relu_s1 = nn.ReLU(inplace=True) def forward(self, x): return self.layer(x)def upsample_like(src, tar): src = F.interpolate(src, size=tar.shape[2:], mode='bilinear') return srcclass DownSample(nn.Module): def __init__(self, ): super(DownSample, self).__init__() self.layer = nn.MaxPool2d(kernel_size=2, stride=2) def forward(self, x): return self.layer(x)class UNet1(nn.Module): def __init__(self, in_channels, mid_channels, out_channels): super(UNet1, self).__init__() self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1) self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1) self.down1 = DownSample() self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down2 = DownSample() self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down3 = DownSample() self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down4 = DownSample() self.conv5 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down5 = DownSample() self.conv6 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.conv7 = ConvolutionLayer(mid_channels, mid_channels, dilation=2) self.conv8 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv9 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv10 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv11 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv12 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv13 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1) def forward(self, x): x0 = self.conv0(x) x1 = self.conv1(x0) d1 = self.down1(x1) x2 = self.conv2(d1) d2 = self.down2(x2) x3 = self.conv3(d2) d3 = self.down3(x3) x4 = self.conv4(d3) d4 = self.down4(x4) x5 = self.conv5(d4) d5 = self.down5(x5) x6 = self.conv6(d5) x7 = self.conv7(x6) x8 = self.conv8(torch.cat((x7, x6), 1)) up1 = upsample_like(x8, x5) x9 = self.conv9(torch.cat((up1, x5), 1)) up2 = upsample_like(x9, x4) x10 = self.conv10(torch.cat((up2, x4), 1)) up3 = upsample_like(x10, x3) x11 = self.conv11(torch.cat((up3, x3), 1)) up4 = upsample_like(x11, x2) x12 = self.conv12(torch.cat((up4, x2), 1)) up5 = upsample_like(x12, x1) x13 = self.conv13(torch.cat((up5, x1), 1)) return x13 + x0class UNet2(nn.Module): def __init__(self, in_channels, mid_channels, out_channels): super(UNet2, self).__init__() self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1) self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1) self.down1 = DownSample() self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down2 = DownSample() self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down3 = DownSample() self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down4 = DownSample() self.conv5 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.conv6 = ConvolutionLayer(mid_channels, mid_channels, dilation=2) self.conv7 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv8 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv9 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv10 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv11 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1) def forward(self, x): x0 = self.conv0(x) x1 = self.conv1(x0) d1 = self.down1(x1) x2 = self.conv2(d1) d2 = self.down2(x2) x3 = self.conv3(d2) d3 = self.down3(x3) x4 = self.conv4(d3) d4 = self.down4(x4) x5 = self.conv5(d4) x6 = self.conv6(x5) x7 = self.conv7(torch.cat((x6, x5), dim=1)) up1 = upsample_like(x7, x4) x8 = self.conv8(torch.cat((up1, x4), dim=1)) up2 = upsample_like(x8, x3) x9 = self.conv9(torch.cat((up2, x3), dim=1)) up3 = upsample_like(x9, x2) x10 = self.conv10(torch.cat((up3, x2), dim=1)) up4 = upsample_like(x10, x1) x11 = self.conv11(torch.cat((up4, x1), dim=1)) return x11 + x0class UNet3(nn.Module): def __init__(self, in_channels, mid_channels, out_channels): super(UNet3, self).__init__() self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1) self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1) self.down1 = DownSample() self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down2 = DownSample() self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down3 = DownSample() self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.conv5 = ConvolutionLayer(mid_channels, mid_channels, dilation=2) self.conv6 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv7 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv8 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv9 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1) def forward(self, x): x0 = self.conv0(x) x1 = self.conv1(x0) d1 = self.down1(x1) x2 = self.conv2(d1) d2 = self.down2(x2) x3 = self.conv3(d2) d3 = self.down3(x3) x4 = self.conv4(d3) x5 = self.conv5(x4) x6 = self.conv6(torch.cat((x5, x4), 1)) up1 = upsample_like(x6, x3) x7 = self.conv7(torch.cat((up1, x3), 1)) up2 = upsample_like(x7, x2) x8 = self.conv8(torch.cat((up2, x2), 1)) up3 = upsample_like(x8, x1) x9 = self.conv9(torch.cat((up3, x1), 1)) return x9 + x0class UNet4(nn.Module): def __init__(self, in_channels, mid_channels=12, out_channels): super(UNet4, self).__init__() self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1) self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1) self.down1 = DownSample() self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.down2 = DownSample() self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1) self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=2) self.conv5 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv6 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1) self.conv7 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1) def forward(self, x): """encode""" x0 = self.conv0(x) x1 = self.conv1(x0) d1 = self.down1(x1) x2 = self.conv2(d1) d2 = self.down2(x2) x3 = self.conv3(d2) x4 = self.conv4(x3) """decode""" x5 = self.conv5(torch.cat((x4, x3), 1)) up1 = upsample_like(x5, x2) x6 = self.conv6(torch.cat((up1, x2), 1)) up2 = upsample_like(x6, x1) x7 = self.conv7(torch.cat((up2, x1), 1)) return x7 + x0class UNet5(nn.Module): def __init__(self, in_channels, mid_channels, out_channels): super(UNet5, self).__init__() self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1) self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1) self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=2) self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=4) self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=8) self.conv5 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=4) self.conv6 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=2) self.conv7 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1) def forward(self, x): x0 = self.conv0(x) x1 = self.conv1(x0) x2 = self.conv2(x1) x3 = self.conv3(x2) x4 = self.conv4(x3) x5 = self.conv5(torch.cat((x4, x3), 1)) x6 = self.conv6(torch.cat((x5, x2), 1)) x7 = self.conv7(torch.cat((x6, x1), 1)) return x7 + x0class U2Net(nn.Module): def __init__(self, in_channels=3, out_channels=1): super(U2Net, self).__init__() self.en_1 = UNet1(in_channels, 32, 64) self.down1 = DownSample() self.en_2 = UNet2(64, 32, 128) self.down2 = DownSample() self.en_3 = UNet3(128, 64, 256) self.down3 = DownSample() self.en_4 = UNet4(256, 128, 512) self.down4 = DownSample() self.en_5 = UNet5(512, 256, 512) self.down5 = DownSample() self.en_6 = UNet5(512, 256, 512) # decoder self.de_5 = UNet5(1024, 256, 512) self.de_4 = UNet4(1024, 128, 256) self.de_3 = UNet3(512, 64, 128) self.de_2 = UNet2(256, 32, 64) self.de_1 = UNet1(128, 16, 64) self.side1 = nn.Conv2d(64, out_channels, kernel_size=(3, 3), padding=1) self.side2 = nn.Conv2d(64, out_channels, kernel_size=(3, 3), padding=1) self.side3 = nn.Conv2d(128, out_channels, kernel_size=(3, 3), padding=1) self.side4 = nn.Conv2d(256, out_channels, kernel_size=(3, 3), padding=1) self.side5 = nn.Conv2d(512, out_channels, kernel_size=(3, 3), padding=1) self.side6 = nn.Conv2d(512, out_channels, kernel_size=(3, 3), padding=1) self.out_conv = nn.Conv2d(6, out_channels, kernel_size=(1, 1)) def forward(self, x): # ------encode ------ x1 = self.en_1(x) d1 = self.down1(x1) x2 = self.en_2(d1) d2 = self.down2(x2) x3 = self.en_3(d2) d3 = self.down3(x3) x4 = self.en_4(d3) d4 = self.down4(x4) x5 = self.en_5(d4) d5 = self.down5(x5) x6 = self.en_6(d5) up1 = upsample_like(x6, x5) # ------decode ------ x7 = self.de_5(torch.cat((up1, x5), dim=1)) up2 = upsample_like(x7, x4) x8 = self.de_4(torch.cat((up2, x4), dim=1)) up3 = upsample_like(x8, x3) x9 = self.de_3(torch.cat((up3, x3), dim=1)) up4 = upsample_like(x9, x2) x10 = self.de_2(torch.cat((up4, x2), dim=1)) up5 = upsample_like(x10, x1) x11 = self.de_1(torch.cat((up5, x1), dim=1)) # side output sup1 = self.side1(x11) sup2 = self.side2(x10) sup2 = upsample_like(sup2, sup1) sup3 = self.side3(x9) sup3 = upsample_like(sup3, sup1) sup4 = self.side4(x8) sup4 = upsample_like(sup4, sup1) sup5 = self.side5(x7) sup5 = upsample_like(sup5, sup1) sup6 = self.side6(x6) sup6 = upsample_like(sup6, sup1) sup0 = self.out_conv(torch.cat((sup1, sup2, sup3, sup4, sup5, sup6), 1)) return torch.sigmoid(sup0)if __name__ == '__main__': u2net = U2Net(3, 1) x = torch.randn(1,3, 512, 512) print(u2net(x).shape) x5 = self.en_5(d4) d5 = self.down5(x5) x6 = self.en_6(d5) up1 = upsample_like(x6, x5) # ------decode ------ x7 = self.de_5(torch.cat((up1, x5), dim=1)) up2 = upsample_like(x7, x4) x8 = self.de_4(torch.cat((up2, x4), dim=1)) up3 = upsample_like(x8, x3) x9 = self.de_3(torch.cat((up3, x3), dim=1)) up4 = upsample_like(x9, x2) x10 = self.de_2(torch.cat((up4, x2), dim=1)) up5 = upsample_like(x10, x1) x11 = self.de_1(torch.cat((up5, x1), dim=1)) # side output sup1 = self.side1(x11) sup2 = self.side2(x10) sup2 = upsample_like(sup2, sup1) sup3 = self.side3(x9) sup3 = upsample_like(sup3, sup1) sup4 = self.side4(x8) sup4 = upsample_like(sup4, sup1) sup5 = self.side5(x7) sup5 = upsample_like(sup5, sup1) sup6 = self.side6(x6) sup6 = upsample_like(sup6, sup1) sup0 = self.out_conv(torch.cat((sup1, sup2, sup3, sup4, sup5, sup6), 1)) return torch.sigmoid(sup0)if __name__ == '__main__': u2net = U2Net(3, 1) x = torch.randn(1,3, 512, 512) print(u2net(x).shape)