深圳幻海软件技术有限公司 欢迎您！
更多

语义分割系列6-Unet++（pytorch实现）

2023-02-28

self nn unet
目录Unet++网络Denseconnectiondeepsupervision模型复现Unet++数据集准备模型训练训练结果Unet++：《UNet++:ANestedU-NetArchitectureforMedicalImageSegmentation》作者对Unet和Unet++的理解：研习

目录

Unet++网络

Dense connection

deep supervision

模型复现

Unet++

数据集准备

模型训练

训练结果

Unet++：《UNet++: A Nested U-Net Architecture for Medical Image Segmentation》

作者对Unet和Unet++的理解：研习U-Net 

延续前文：语义分割系列2-Unet（pytorch实现）

本文将介绍Unet++网络，在pytorch框架上复现Unet++，并在Camvid数据集上进行训练。

Unet++网络

Dense connection

Unet++继承了Unet的结构，同时又借鉴了DenseNet的稠密连接方式（图1中各种分支）。

图1 Unet++网络结构

作者通过各层之间的稠密连接，互相连接起来，就像Denset那样，前前后后每一个模块互相作用，每一个模块都能看到彼此，那对彼此互相熟悉，分割效果自然就会变好。

在实际分割中，一次次的下采样自然会丢掉一些细节特征，在Unet中是使用skip connection来恢复这些细节，但能否做的更好呢？Unet++就给出了答案，这种稠密连接的方式，每一层都尽量多的保存这种细节信息和全局信息，一层层之间架起桥梁互相沟通，最后共享给最后一层，实现全局信息和局部信息的保留和重构。 

deep supervision

当然，简单的将各个模块连接起来是会实现很好的效果。而我们又能发现，一个Unet++其实是很多个不同深度的Unet++叠加。那么，每一个深度的Unet++是不是就都可以输出一个loss？答案自然是可以的。

所以，作者提出了deep supervision，也就是监督每一个深度的Unet++的输出，通过一定的方式来叠加Loss（比如加权的方式），这样就得到了一个经由1、2、3、4层的Unet++的加权Loss（图2 不同深度Unet++融合）。

图2 不同深度Unet++融合

那么，deep supervision又有什么用呢？-剪枝

既然Unet++由多个不同深度的Unet++叠加起来，那么随意去掉一层，前向传播的梯度不会受到任何变化，但你发现Unet++第三个输出的效果和第四个输出效果差不多时，那就可以毫不犹豫删去4层深度的Unet++。比如，直接删去图3中棕色部分，就可以实现剪枝。这样，就得到了更加轻量化的网络。

图3 剪枝模型

模型复现

Unet++

为了更直观一些，我把代码中的所有符号都和网络结构中对应上了。

  1. import torch
  2. import torch.nn as nn
  3. class ContinusParalleConv(nn.Module):
  4.     # 一个连续的卷积模块，包含BatchNorm 在前 和 在后 两种模式
  5.     def __init__(self, in_channels, out_channels, pre_Batch_Norm = True):
  6.         super(ContinusParalleConv, self).__init__()
  7.         self.in_channels = in_channels
  8.         self.out_channels = out_channels
  9.  
  10.         if pre_Batch_Norm:
  11.           self.Conv_forward = nn.Sequential(
  12.             nn.BatchNorm2d(self.in_channels),
  13.             nn.ReLU(),
  14.             nn.Conv2d(self.in_channels, self.out_channels, 3, padding=1),
  15.             nn.BatchNorm2d(out_channels),
  16.             nn.ReLU(),
  17.             nn.Conv2d(self.out_channels, self.out_channels, 3, padding=1))
  18.  
  19.         else:
  20.           self.Conv_forward = nn.Sequential(
  21.             nn.Conv2d(self.in_channels, self.out_channels, 3, padding=1),
  22.             nn.BatchNorm2d(out_channels),
  23.             nn.ReLU(),
  24.             nn.Conv2d(self.out_channels, self.out_channels, 3, padding=1),
  25.             nn.BatchNorm2d(self.out_channels),
  26.             nn.ReLU())
  27.  
  28.     def forward(self, x):
  29.         x = self.Conv_forward(x)
  30.         return x
  31.  
  32. class UnetPlusPlus(nn.Module):
  33.     def __init__(self, num_classes, deep_supervision=False):
  34.         super(UnetPlusPlus, self).__init__()
  35.         self.num_classes = num_classes
  36.         self.deep_supervision = deep_supervision
  37.         self.filters = [64, 128, 256, 512, 1024]
  38.         
  39.         self.CONV3_1 = ContinusParalleConv(512*2, 512, pre_Batch_Norm = True)
  40.  
  41.         self.CONV2_2 = ContinusParalleConv(256*3, 256, pre_Batch_Norm = True)
  42.         self.CONV2_1 = ContinusParalleConv(256*2, 256, pre_Batch_Norm = True)
  43.  
  44.         self.CONV1_1 = ContinusParalleConv(128*2, 128, pre_Batch_Norm = True)
  45.         self.CONV1_2 = ContinusParalleConv(128*3, 128, pre_Batch_Norm = True)
  46.         self.CONV1_3 = ContinusParalleConv(128*4, 128, pre_Batch_Norm = True)
  47.  
  48.         self.CONV0_1 = ContinusParalleConv(64*2, 64, pre_Batch_Norm = True)
  49.         self.CONV0_2 = ContinusParalleConv(64*3, 64, pre_Batch_Norm = True)
  50.         self.CONV0_3 = ContinusParalleConv(64*4, 64, pre_Batch_Norm = True)
  51.         self.CONV0_4 = ContinusParalleConv(64*5, 64, pre_Batch_Norm = True)
  52.  
  53.  
  54.         self.stage_0 = ContinusParalleConv(3, 64, pre_Batch_Norm = False)
  55.         self.stage_1 = ContinusParalleConv(64, 128, pre_Batch_Norm = False)
  56.         self.stage_2 = ContinusParalleConv(128, 256, pre_Batch_Norm = False)
  57.         self.stage_3 = ContinusParalleConv(256, 512, pre_Batch_Norm = False)
  58.         self.stage_4 = ContinusParalleConv(512, 1024, pre_Batch_Norm = False)
  59.  
  60.         self.pool = nn.MaxPool2d(2)
  61.     
  62.         self.upsample_3_1 = nn.ConvTranspose2d(in_channels=1024, out_channels=512, kernel_size=4, stride=2, padding=1) 
  63.  
  64.         self.upsample_2_1 = nn.ConvTranspose2d(in_channels=512, out_channels=256, kernel_size=4, stride=2, padding=1) 
  65.         self.upsample_2_2 = nn.ConvTranspose2d(in_channels=512, out_channels=256, kernel_size=4, stride=2, padding=1) 
  66.  
  67.         self.upsample_1_1 = nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=4, stride=2, padding=1) 
  68.         self.upsample_1_2 = nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=4, stride=2, padding=1) 
  69.         self.upsample_1_3 = nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=4, stride=2, padding=1) 
  70.  
  71.         self.upsample_0_1 = nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=4, stride=2, padding=1) 
  72.         self.upsample_0_2 = nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=4, stride=2, padding=1) 
  73.         self.upsample_0_3 = nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=4, stride=2, padding=1) 
  74.         self.upsample_0_4 = nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=4, stride=2, padding=1) 
  75.  
  76.         
  77.         # 分割头
  78.         self.final_super_0_1 = nn.Sequential(
  79.           nn.BatchNorm2d(64),
  80.           nn.ReLU(),
  81.           nn.Conv2d(64, self.num_classes, 3, padding=1),
  82.         )        
  83.         self.final_super_0_2 = nn.Sequential(
  84.           nn.BatchNorm2d(64),
  85.           nn.ReLU(),
  86.           nn.Conv2d(64, self.num_classes, 3, padding=1),
  87.         )        
  88.         self.final_super_0_3 = nn.Sequential(
  89.           nn.BatchNorm2d(64),
  90.           nn.ReLU(),
  91.           nn.Conv2d(64, self.num_classes, 3, padding=1),
  92.         )        
  93.         self.final_super_0_4 = nn.Sequential(
  94.           nn.BatchNorm2d(64),
  95.           nn.ReLU(),
  96.           nn.Conv2d(64, self.num_classes, 3, padding=1),
  97.         )        
  98.  
  99.         
  100.     def forward(self, x):
  101.         x_0_0 = self.stage_0(x)
  102.         x_1_0 = self.stage_1(self.pool(x_0_0))
  103.         x_2_0 = self.stage_2(self.pool(x_1_0))
  104.         x_3_0 = self.stage_3(self.pool(x_2_0))
  105.         x_4_0 = self.stage_4(self.pool(x_3_0))
  106.         
  107.         x_0_1 = torch.cat([self.upsample_0_1(x_1_0) , x_0_0], 1)
  108.         x_0_1 =  self.CONV0_1(x_0_1)
  109.         
  110.         x_1_1 = torch.cat([self.upsample_1_1(x_2_0), x_1_0], 1)
  111.         x_1_1 = self.CONV1_1(x_1_1)
  112.         
  113.         x_2_1 = torch.cat([self.upsample_2_1(x_3_0), x_2_0], 1)
  114.         x_2_1 = self.CONV2_1(x_2_1)
  115.         
  116.         x_3_1 = torch.cat([self.upsample_3_1(x_4_0), x_3_0], 1)
  117.         x_3_1 = self.CONV3_1(x_3_1)
  118.  
  119.         x_2_2 = torch.cat([self.upsample_2_2(x_3_1), x_2_0, x_2_1], 1)
  120.         x_2_2 = self.CONV2_2(x_2_2)
  121.         
  122.         x_1_2 = torch.cat([self.upsample_1_2(x_2_1), x_1_0, x_1_1], 1)
  123.         x_1_2 = self.CONV1_2(x_1_2)
  124.         
  125.         x_1_3 = torch.cat([self.upsample_1_3(x_2_2), x_1_0, x_1_1, x_1_2], 1)
  126.         x_1_3 = self.CONV1_3(x_1_3)
  127.  
  128.         x_0_2 = torch.cat([self.upsample_0_2(x_1_1), x_0_0, x_0_1], 1)
  129.         x_0_2 = self.CONV0_2(x_0_2)
  130.         
  131.         x_0_3 = torch.cat([self.upsample_0_3(x_1_2), x_0_0, x_0_1, x_0_2], 1)
  132.         x_0_3 = self.CONV0_3(x_0_3)
  133.         
  134.         x_0_4 = torch.cat([self.upsample_0_4(x_1_3), x_0_0, x_0_1, x_0_2, x_0_3], 1)
  135.         x_0_4 = self.CONV0_4(x_0_4)
  136.     
  137.     
  138.         if self.deep_supervision:
  139.             out_put1 = self.final_super_0_1(x_0_1)
  140.             out_put2 = self.final_super_0_2(x_0_2)
  141.             out_put3 = self.final_super_0_3(x_0_3)
  142.             out_put4 = self.final_super_0_4(x_0_4)
  143.             return [out_put1, out_put2, out_put3, out_put4]
  144.         else:
  145.             return self.final_super_0_4(x_0_4)
  146.  
  147.  
  148. if __name__ == "__main__":
  149.     print("deep_supervision: False")
  150.     deep_supervision = False
  151.     device = torch.device('cpu')
  152.     inputs = torch.randn((1, 3, 224, 224)).to(device)
  153.     model = UnetPlusPlus(num_classes=3, deep_supervision=deep_supervision).to(device)
  154.     outputs = model(inputs)
  155.     print(outputs.shape)    
  156.     
  157.     print("deep_supervision: True")
  158.     deep_supervision = True
  159.     model = UnetPlusPlus(num_classes=3, deep_supervision=deep_supervision).to(device)
  160.     outputs = model(inputs)
  161.     for out in outputs:
  162.       print(out.shape)
  163.  
  164.

测试结果如下 

数据集准备

数据集使用Camvid数据集，可在CamVid数据集的创建和使用-pytorch中参考构建方法。

  1. # 导入库
  2. import os
  3. os.environ['CUDA_VISIBLE_DEVICES'] = '0'
  4.  
  5. import torch
  6. import torch.nn as nn
  7. import torch.optim as optim
  8. import torch.nn.functional as F
  9. from torch import optim
  10. from torch.utils.data import Dataset, DataLoader, random_split
  11. from tqdm import tqdm
  12. import warnings
  13. warnings.filterwarnings("ignore")
  14. import os.path as osp
  15. import matplotlib.pyplot as plt
  16. from PIL import Image
  17. import numpy as np
  18. import albumentations as A
  19. from albumentations.pytorch.transforms import ToTensorV2
  20.  
  21.  
  22. torch.manual_seed(17)
  23. # 自定义数据集CamVidDataset
  24. class CamVidDataset(torch.utils.data.Dataset):
  25.     """CamVid Dataset. Read images, apply augmentation and preprocessing transformations.
  26.     
  27.     Args:
  28.         images_dir (str): path to images folder
  29.         masks_dir (str): path to segmentation masks folder
  30.         class_values (list): values of classes to extract from segmentation mask
  31.         augmentation (albumentations.Compose): data transfromation pipeline 
  32.             (e.g. flip, scale, etc.)
  33.         preprocessing (albumentations.Compose): data preprocessing 
  34.             (e.g. noralization, shape manipulation, etc.)
  35.     """
  36.     
  37.     def __init__(self, images_dir, masks_dir):
  38.         self.transform = A.Compose([
  39.             A.Resize(224, 224),
  40.             A.HorizontalFlip(),
  41.             A.VerticalFlip(),
  42.             A.Normalize(),
  43.             ToTensorV2(),
  44.         ]) 
  45.         self.ids = os.listdir(images_dir)
  46.         self.images_fps = [os.path.join(images_dir, image_id) for image_id in self.ids]
  47.         self.masks_fps = [os.path.join(masks_dir, image_id) for image_id in self.ids]
  48.  
  49.     
  50.     def __getitem__(self, i):
  51.         # read data
  52.         image = np.array(Image.open(self.images_fps[i]).convert('RGB'))
  53.         mask = np.array( Image.open(self.masks_fps[i]).convert('RGB'))
  54.         image = self.transform(image=image,mask=mask)
  55.         
  56.         return image['image'], image['mask'][:,:,0]
  57.         
  58.     def __len__(self):
  59.         return len(self.ids)
  60.     
  61.     
  62. # 设置数据集路径
  63. DATA_DIR = r'dataset\camvid' # 根据自己的路径来设置
  64. x_train_dir = os.path.join(DATA_DIR, 'train_images')
  65. y_train_dir = os.path.join(DATA_DIR, 'train_labels')
  66. x_valid_dir = os.path.join(DATA_DIR, 'valid_images')
  67. y_valid_dir = os.path.join(DATA_DIR, 'valid_labels')
  68.     
  69. train_dataset = CamVidDataset(
  70.     x_train_dir, 
  71.     y_train_dir, 
  72. )
  73. val_dataset = CamVidDataset(
  74.     x_valid_dir, 
  75.     y_valid_dir, 
  76. )
  77.  
  78. train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True,drop_last=True)
  79. val_loader = DataLoader(val_dataset, batch_size=8, shuffle=True,drop_last=True)

模型训练

  1. model = UnetPlusPlus(num_classes=33).cuda()
  2. #载入预训练模型
  3. #model.load_state_dict(torch.load(r"checkpoints/Unet++_25.pth"),strict=False)
  1. from d2l import torch as d2l
  2. from tqdm import tqdm
  3. import pandas as pd
  4. #损失函数选用多分类交叉熵损失函数
  5. lossf = nn.CrossEntropyLoss(ignore_index=255)
  6. #选用adam优化器来训练
  7. optimizer = optim.SGD(model.parameters(),lr=0.1)
  8. scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1, last_epoch=-1)
  9.  
  10. #训练50轮
  11. epochs_num = 50
  12. def train_ch13(net, train_iter, test_iter, loss, trainer, num_epochs,scheduler,
  13.                devices=d2l.try_all_gpus()):
  14.     timer, num_batches = d2l.Timer(), len(train_iter)
  15.     animator = d2l.Animator(xlabel='epoch', xlim=[1, num_epochs], ylim=[0, 1],
  16.                             legend=['train loss', 'train acc', 'test acc'])
  17.     net = nn.DataParallel(net, device_ids=devices).to(devices[0])
  18.     
  19.     loss_list = []
  20.     train_acc_list = []
  21.     test_acc_list = []
  22.     epochs_list = []
  23.     time_list = []
  24.     for epoch in range(num_epochs):
  25.         # Sum of training loss, sum of training accuracy, no. of examples,
  26.         # no. of predictions
  27.         metric = d2l.Accumulator(4)
  28.         for i, (features, labels) in enumerate(train_iter):
  29.             timer.start()
  30.             l, acc = d2l.train_batch_ch13(
  31.                 net, features, labels.long(), loss, trainer, devices)
  32.             metric.add(l, acc, labels.shape[0], labels.numel())
  33.             timer.stop()
  34.             if (i + 1) % (num_batches // 5) == 0 or i == num_batches - 1:
  35.                 animator.add(epoch + (i + 1) / num_batches,
  36.                              (metric[0] / metric[2], metric[1] / metric[3],
  37.                               None))
  38.         test_acc = d2l.evaluate_accuracy_gpu(net, test_iter)
  39.         animator.add(epoch + 1, (None, None, test_acc))
  40.         scheduler.step()
  41.  
  42.         print(f"epoch {epoch+1} --- loss {metric[0] / metric[2]:.3f} ---  train acc {metric[1] / metric[3]:.3f} --- test acc {test_acc:.3f} --- cost time {timer.sum()}")
  43.         
  44.         #---------保存训练数据---------------
  45.         df = pd.DataFrame()
  46.         loss_list.append(metric[0] / metric[2])
  47.         train_acc_list.append(metric[1] / metric[3])
  48.         test_acc_list.append(test_acc)
  49.         epochs_list.append(epoch)
  50.         time_list.append(timer.sum())
  51.         
  52.         df['epoch'] = epochs_list
  53.         df['loss'] = loss_list
  54.         df['train_acc'] = train_acc_list
  55.         df['test_acc'] = test_acc_list
  56.         df['time'] = time_list
  57.         df.to_excel("savefile/Unet++_camvid1.xlsx")
  58.         #----------------保存模型-------------------
  59.         if np.mod(epoch+1, 5) == 0:
  60.             torch.save(model.state_dict(), f'checkpoints/Unet++_{epoch+1}.pth')

开始训练

train_ch13(model, train_loader, val_loader, lossf, optimizer, epochs_num,scheduler)

训练结果

Unet++训练结果
文章知识点与官方知识档案匹配，可进一步学习相关知识
OpenCV技能树OpenCV中的深度学习图像分类13287 人正在系统学习中
ACM基础(一)
国产数据库，又有好消息！