0
0

语义分割系列7-Attention Unet（pytorch实现）

继前文Unet和Unet++之后，本文将介绍Attention Unet。

Attention Unet地址，《Attention U-Net: Learning Where to Look for the Pancreas》。

AttentionUnet

Attention Unet发布于2018年，主要应用于医学领域的图像分割，全文中主要以肝脏的分割论证。

论文中心

Attention Unet主要的中心思想就是提出来Attention gate模块，使用soft-attention替代hard-attention，将attention集成到Unet的跳跃连接和上采样模块中，实现空间上的注意力机制。通过attention机制来抑制图像中的无关信息，突出局部的重要特征。

网络架构


图1 AttentionUnet模型

Attention Unet的模型结构和Unet十分相像，只是增加了Attention Gate模块来对skip connection和upsampling层做attention机制（图2）。


图2 Attention Gate模块

在Attention Gate模块中，g和xl分别为skip connection的输出和下一层的输出，如图3。


图3 Attention Gate的输入

需要注意的是，在计算Wg和Wx后，对两者进行相加。但是，此时g的维度和xl的维度并不相等，则需要对g做下采样或对xl做上采样。（我倾向于对xl做上采样，因为在原本的Unet中，在Decoder就需要对下一层做上采样，所以，直接使用这个上采样结果可以减少网络计算）。

Wg和Wx经过相加，ReLU激活，1x1x1卷积，Sigmoid激活，生成一个权重信息，将这个权重与原始输入xl相乘，得到了对xl的attention激活。这就是Attenton Gate的思想。

Attenton Gate还有一个比较重要的特点是：这个权重可以经由网络学习！因为soft-attention是可微的，可以微分的attention就可以通过神经网络算出梯度并且前向传播和后向反馈来学习得到attention的权重。以此来学习更重要的特征。

模型复现

Attention Unet代码

  1. import torch
  2. import torch.nn as nn
  3. #Attention gate代码
  4. class AttentionBlock(nn.Module):
  5.     def __init__(self, F_g, F_l, F_int):
  6.         super(AttentionBlock, self).__init__()
  7.         self.W_g = nn.Sequential(
  8.             nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=False),
  9.             nn.BatchNorm2d(F_int)
  10.         )
  11.         
  12.         self.W_x = nn.Sequential(
  13.             nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=False),
  14.             nn.BatchNorm2d(F_int)
  15.         )
  16.         
  17.         self.psi = nn.Sequential(
  18.             nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=False),
  19.             nn.BatchNorm2d(1),
  20.             nn.Sigmoid()
  21.         )
  22.         
  23.         self.relu = nn.ReLU(inplace=True)
  24.         
  25.     def forward(self, g, x):
  26.         g = self.W_g(g)
  27.         x = self.W_x(x)
  28.         psi = self.relu(g+x)
  29.         psi = self.psi(psi)
  30.         
  31.         return x*psi
  32.         
  33.         
  34.         
  35. #AttentionUnet代码
  36. class AttentionUnet(nn.Module):
  37.     def __init__(self, num_classes):
  38.         super(AttentionUnet, self).__init__()
  39.         self.stage_1 = nn.Sequential(
  40.             nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3,padding=1),
  41.             nn.BatchNorm2d(32),
  42.             nn.ReLU(),
  43.             nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3,padding=1),
  44.             nn.BatchNorm2d(64),
  45.             nn.ReLU(),
  46.             nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3,padding=1),
  47.             nn.BatchNorm2d(64),
  48.             nn.ReLU(),
  49.         )
  50.         
  51.         self.stage_2 = nn.Sequential(
  52.             nn.MaxPool2d(kernel_size=2),
  53.             nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3,padding=1),
  54.             nn.BatchNorm2d(128),
  55.             nn.ReLU(),
  56.             nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3,padding=1),
  57.             nn.BatchNorm2d(128),
  58.             nn.ReLU(),
  59.         )
  60.         
  61.         self.stage_3 = nn.Sequential(
  62.             nn.MaxPool2d(kernel_size=2),
  63.             nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3,padding=1),
  64.             nn.BatchNorm2d(256),
  65.             nn.ReLU(),
  66.             nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3,padding=1),
  67.             nn.BatchNorm2d(256),
  68.             nn.ReLU(),
  69.         )
  70.         
  71.         self.stage_4 = nn.Sequential(
  72.             nn.MaxPool2d(kernel_size=2),
  73.             nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3,padding=1),
  74.             nn.BatchNorm2d(512),
  75.             nn.ReLU(),
  76.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3,padding=1),
  77.             nn.BatchNorm2d(512),
  78.             nn.ReLU(),
  79.         )
  80.         
  81.         self.stage_5 = nn.Sequential(
  82.             nn.MaxPool2d(kernel_size=2),
  83.             nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=3,padding=1),
  84.             nn.BatchNorm2d(1024),
  85.             nn.ReLU(),
  86.             nn.Conv2d(in_channels=1024, out_channels=1024, kernel_size=3,padding=1),
  87.             nn.BatchNorm2d(1024),
  88.             nn.ReLU(),
  89.         )
  90.         
  91.         self.upsample_4 = nn.Sequential(
  92.             nn.ConvTranspose2d(in_channels=1024, out_channels=512,kernel_size=4,stride=2, padding=1) 
  93.         )
  94.         self.upsample_3 = nn.Sequential(
  95.             nn.ConvTranspose2d(in_channels=512, out_channels=256,kernel_size=4,stride=2, padding=1) 
  96.         )
  97.         self.upsample_2 = nn.Sequential(
  98.             nn.ConvTranspose2d(in_channels=256, out_channels=128,kernel_size=4,stride=2, padding=1) 
  99.         )
  100.         self.upsample_1 = nn.Sequential(
  101.             nn.ConvTranspose2d(in_channels=128, out_channels=64,kernel_size=4,stride=2, padding=1) 
  102.         )
  103.         
  104.         self.stage_up_4 = nn.Sequential(
  105.             nn.Conv2d(in_channels=1024, out_channels=512, kernel_size=3,padding=1),
  106.             nn.BatchNorm2d(512),
  107.             nn.ReLU(),
  108.             nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3,padding=1),
  109.             nn.BatchNorm2d(512),
  110.             nn.ReLU()
  111.         )
  112.         
  113.         self.stage_up_3 = nn.Sequential(
  114.             nn.Conv2d(in_channels=512, out_channels=256, kernel_size=3,padding=1),
  115.             nn.BatchNorm2d(256),
  116.             nn.ReLU(),
  117.             nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3,padding=1),
  118.             nn.BatchNorm2d(256),
  119.             nn.ReLU()
  120.         )    
  121.         
  122.         self.stage_up_2 = nn.Sequential(
  123.             nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3,padding=1),
  124.             nn.BatchNorm2d(128),
  125.             nn.ReLU(),
  126.             nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3,padding=1),
  127.             nn.BatchNorm2d(128),
  128.             nn.ReLU()
  129.         ) 
  130.         self.stage_up_1 = nn.Sequential(
  131.             nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3,padding=1),
  132.             nn.BatchNorm2d(64),
  133.             nn.ReLU(),
  134.             nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3,padding=1),
  135.             nn.BatchNorm2d(64),
  136.             nn.ReLU()
  137.         )    
  138.         
  139.         self.Attentiongate1 = AttentionBlock(512, 512, 512)
  140.         self.Attentiongate2 = AttentionBlock(256, 256, 256)
  141.         self.Attentiongate3 = AttentionBlock(128, 128, 128)
  142.         
  143.         self.final = nn.Sequential(
  144.             nn.Conv2d(in_channels=64, out_channels=num_classes, kernel_size=3, padding=1),
  145.         )
  147.         
  148.     def forward(self, x):
  149.         x = x.float()
  150.         #下采样过程
  151.         stage_1 = self.stage_1(x)
  152.         stage_2 = self.stage_2(stage_1)
  153.         stage_3 = self.stage_3(stage_2)
  154.         stage_4 = self.stage_4(stage_3)
  155.         stage_5 = self.stage_5(stage_4)
  156.         
  157.  
  159.         up_4 = self.upsample_4(stage_5)
  160.         stage_4 = self.Attentiongate1(up_4, stage_4)
  161.         up_4_conv = self.stage_up_4(torch.cat([up_4, stage_4], dim=1))
  162.         
  163.         
  164.         up_3 = self.upsample_3(up_4_conv)
  165.         stage_3 = self.Attentiongate2(up_3, stage_3)
  166.         up_3_conv = self.stage_up_3(torch.cat([up_3, stage_3], dim=1))
  167.         
  168.         up_2 = self.upsample_2(up_3_conv)
  169.         stage_2 = self.Attentiongate3(up_2, stage_2)
  170.         up_2_conv = self.stage_up_2(torch.cat([up_2, stage_2], dim=1))
  171.         
  172.         up_1 = self.upsample_1(up_2_conv)
  173.         up_1_conv = self.stage_up_1(torch.cat([up_1, stage_1], dim=1))   
  174.         
  175.         output = self.final(up_1_conv)
  176.         
  177.         return output

数据集

数据集依旧使用Camvid数据集，见Camvid数据集的构建和使用。

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

模型训练

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

开始训练

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

训练结果

标签： pytorch 深度学习 人工智能
本文转载自: https://blog.csdn.net/yumaomi/article/details/124866235
版权归原作者 yumaomi 所有， 如有侵权，请联系我们删除。
登录后发布评论

“语义分割系列7-Attention Unet（pytorch实现）”的评论:

还没有评论
关于作者
...
overfit同步小助手
文章同步
相关阅读

网络安全法-网络运行安全

使用selenium/drissionpage时如何阻止chrome自动跳转http到https

docker desktop 里部署的Open WebUI 管理员密码忘记了的处理方法

在ubuntu20.04中搭建onsite比赛运行环境

利用开源的低代码表单设计器FcDesigner高效管理和渲染复杂表单结构

Kafka学习笔记

【前端】浏览器输入url到页面呈现发生了什么？