动手实现基于pytorch框架的UNet模型

最新推荐文章于 2024-04-30 22:45:00 发布
最新推荐文章于 2024-04-30 22:45:00 发布
阅读量1.8k 收藏 26
点赞数 1
分类专栏： DeepLearning 文章标签： pytorch 深度学习 自动驾驶 UNet segmention
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/BXD1314/article/details/120386732
版权

前言

最近在学习CNN 图像分割相关内容，接触到了UNet 网络，UNet是一个很经典的网络，因其结构像字母U得名，对于一般的图像分割有显著的效果。UNet的网络结构是一个U形结构，左半边是Encoder，右半边是Decoder。Encode部分，下采样不断的增大channel，宽高减半，并提取图像的特征，但是丢弃了图像的位置信息。Decoder 上采样，upconvolution，融合下采样的图像特征并恢复图像的位置信息

  • UNet 结构图
    在这里插入图片描述
  • 关于跟详细的实现内容可以阅读 UNet论文

1.首先，图中的灰色箭头（copy and crop）目的是将浅层特征与深层特征融合，这样可以既保留浅层特征图中较高精度的特征信息，也可以利用深层特征图中抽象的语义信息。
2.其次，在下采样过程中，特征图缩小的尺度是上一层的一半；而在上采样过程中特征图变为上一层的一倍。通道数量变化相反。
3.下采样的时候卷积层特征尺度变化小，原论文使用max pooling进行尺度缩小；上采样也一样，使用upsampling+conv进行尺度增大。

  • 考虑到max pooling会丢失位置信息，决定使用卷积代替它；
  • 使用转置卷积替代简单的上采样（插值），这样既能实现同样的效果，也能加深网络。

但以上做法有利有弊端，如果网络够深，使用卷积是计算速度变慢而且可能产生过拟合

下面开始动手实现Unet结构

实现UNet

下采样的卷积实现

对与下采样过程就是

  1. Conv 3x3,BN,Relu
  2. Conv 3x3,BN,Relu
  3. max pool 2x2

也就是如下图的意思
在这里插入图片描述
但是为了下采样过程和上采样过程衔接部分代码好实现，我连接顺序是：

  1. max pool 2x2
  2. Conv 3x3,BN,Relu
  3. Conv 3x3,BN,Relu
    也就是如下图的意思：

Alt

class DoubleConv(nn.Module):
    """convolution->BN->Relu"""
    def __init__(self,in_channels,out_channels,mid_channels = None):
        super().__init__()

        #这部分和论文不一样,多了个中间输入的channel
        if not mid_channels:
            mid_channels = out_channels
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels,mid_channels,kernel_size=3,padding=1),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(),
            nn.Conv2d(mid_channels,out_channels,kernel_size=3,padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
        )

    def forward(self,x):
        return self.double_conv(x)


#编码器encoder
class Down(nn.Module):
    """Downscaling with maxpool then double conv"""
    def __init__(self, in_channels,out_channels,maxPool = True):
        super().__init__()

        maxPool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels,out_channels),
        )
        
        #使用卷积代替maxPoling 下采样,这样不会丢失位置信息
        #但是如果网络太深，会产生过拟合
        down_conv = nn.Sequential(
            nn.Conv2d(in_channels,out_channels,kernel_size=3,stride=2,padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            DoubleConv(in_channels,out_channels),
            
        )

        self.downsample = ( maxPool_conv if maxPool  else  down_conv)

      
    def forward(self,x):
        return self.downsample(x)

上采样过程

融合下采样数据特征和恢复尺寸

  • up-conv
  • Conv 3x3,BN,Relu
  • Conv 3x3,BN,Relu
  • copy and crop
    在这里插入图片描述
#解码器Decoder
class Up(nn.Module):
    """Downscaling with maxpool then double conv"""
    
    def __init__(self,in_channels,out_channels, bilinear=True):
        super().__init__()

        #如果是双线性插值，使用普通卷积来减少通道
        if bilinear:
            self.up = Upsample(scale_factor=2,mode='bilinear',align_corners=True)
            self.conv = DoubleConv(in_channels,out_channels,in_channels // 2)
        else:
            #采用转置卷积代替上采样，out_channel 是in_channels的一半
            self.up = nn.ConvTranspose2d(in_channels,in_channels//2,kernel_size=2,stride=2)
            self.conv = DoubleConv(in_channels, out_channels)

    def forward(self,x1,x2):
        #据论文，需要进行融合还原之前的尺寸
        #x1是上采样获得的特征
        #x2是下采样获得的特征
        x1 = self.up(x1)
        if (x1.size(2) != x2.size(2)) or (x1.size(3) != x2.size(3)):
            #input is CHW
            #这个是解决填充不一致的问题
            diffY = x2.size()[2] - x1.size()[2]
            diffX = x2.size()[3] - x1.size()[3]
            #print('sizes',x1.size(),x2.size(),diffX // 2, diffX - diffX//2, diffY // 2, diffY - diffY//2)
            x1 = F.pad(x1, (diffX // 2, diffX - diffX//2,
                            diffY // 2, diffY - diffY//2))
            #print("pad x1:",x1.size())
        x = torch.cat([x2, x1], dim=1)
        x = self.conv(x)
        return x

最后一层输出层定义(out segmentation map)

class outlayer(nn.Module):
    def __init__(self,in_channels,out_channels):
        super().__init__()
        self.conv = nn.Conv2d(in_channels,out_channels,kernel_size=1)
    def forward(self,x):
        x =self.conv(x)
        return x

按步骤组织UNet结构

class UNet(nn.Module):
    def __init__(self,n_classes,n_channels = 3,bilinear=True) -> None:
        super().__init__()
        self.n_channels = n_channels
        self.n_channels = n_classes
        self.bilinear = bilinear

        self.start = DoubleConv(n_channels,64)

        self.down1 = Down(64,128)
        self.down2 = Down(128,256)
        self.down3 = Down(256,512)
        #self.down4 = Down(512,1024)
        factor = 2 if bilinear else 1
        self.down4 = Down(512, 1024 // factor)

        self.up1 = Up(1024, 512 // factor, bilinear)
        self.up2 = Up(512, 256 // factor, bilinear)
        self.up3 = Up(256, 128 // factor, bilinear)
        self.up4 = Up(128, 64, bilinear)
        self.final_conv = outlayer(64, n_classes)

        # self.up1 = Up(1024, 512,bilinear)
        # self.up2 = Up(512, 256,bilinear)
        # self.up3 = Up(256, 128 ,bilinear)
        # self.up4 = Up(128, 64,bilinear)
        # self.final_conv = outlayer(64, n_classes)

    def forward(self,x):
        x0 = self.start(x) #3-64
        #print(x0.shape)
        x1 = self.down1(x0)#64-128
        #print(f"x1.shape:\n{x1.shape}")
        x2 = self.down2(x1)#128-246
        #print(f"x2.shape:\n{x2.shape}")
        x3 = self.down3(x2)#256-512
        #print(f"x3.shape:\n{x3.shape}")
        x4 = self.down4(x3)#512-1024
        #print(f"x4.shape:\n{x4.shape}")

        x = self.up1(x4, x3)#1024-512
        #print(f"x.shape:\n{x.shape}")
        x = self.up2(x, x2)#512-256
        #print(f"x.shape:\n{x.shape}")
        x = self.up3(x, x1)#256-128
        #print(f"x.shape:\n{x.shape}")
        x = self.up4(x, x0)#128-64
        #print(f"x.shape:\n{x.shape}")
        logits = self.final_conv(x)
        #print(f"logits:\n{logits.shape}")
        return logits

测试UNet结构是否正确

if __name__ == '__main__':
    net = UNet(n_channels=1,n_classes=2,bilinear=False)
    dev = ('cuda:0' if torch.cuda.is_available() else 'cpu')
    print(dev)
    x = torch.randn(1,1,572,572)
    out = net(x).to(dev)
    print(net)
    print(out.shape)
    '''输出结果'''
    cuda:0
C:\Users\bxd\AppData\Local\Programs\Python\Python39\lib\site-packages\torch\nn\functional.py:718: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at  ..\c10/core/TensorImpl.h:1156.)
  return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
UNet(
  (start): DoubleConv(
    (double_conv): Sequential(
      (0): Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)     
      (2): ReLU()
      (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)     
      (5): ReLU()
    )
  )
  (down1): Down(
    (downsample): Sequential(
      (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (1): DoubleConv(
        (double_conv): Sequential(
          (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU()
          (3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (4): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (5): ReLU()
        )
      )
    )
  )
  (down2): Down(
    (downsample): Sequential(
      (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (1): DoubleConv(
        (double_conv): Sequential(
          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU()
          (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (5): ReLU()
        )
      )
    )
  )
  (down3): Down(
    (downsample): Sequential(
      (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (1): DoubleConv(
        (double_conv): Sequential(
          (0): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU()
          (3): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (4): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (5): ReLU()
        )
      )
    )
  )
  (down4): Down(
    (downsample): Sequential(
      (0): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (1): DoubleConv(
        (double_conv): Sequential(
          (0): Conv2d(512, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU()
          (3): Conv2d(1024, 1024, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (4): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (5): ReLU()
        )
      )
    )
  )
  (up1): Up(
    (up): ConvTranspose2d(1024, 512, kernel_size=(2, 2), stride=(2, 2))
    (conv): DoubleConv(
      (double_conv): Sequential(
        (0): Conv2d(1024, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): ReLU()
        (3): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (5): ReLU()
      )
    )
  )
  (up2): Up(
    (up): ConvTranspose2d(512, 256, kernel_size=(2, 2), stride=(2, 2))
    (conv): DoubleConv(
      (double_conv): Sequential(
        (0): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): ReLU()
        (3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (5): ReLU()
      )
    )
  )
  (up3): Up(
    (up): ConvTranspose2d(256, 128, kernel_size=(2, 2), stride=(2, 2))
    (conv): DoubleConv(
      (double_conv): Sequential(
        (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): ReLU()
        (3): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (5): ReLU()
      )
    )
  )
  (up4): Up(
    (up): ConvTranspose2d(128, 64, kernel_size=(2, 2), stride=(2, 2))
    (conv): DoubleConv(
      (double_conv): Sequential(
        (0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): ReLU()
        (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (5): ReLU()
      )
    )
  )
  (final_conv): outlayer(
    (conv): Conv2d(64, 2, kernel_size=(1, 1), stride=(1, 1))
  )
)
torch.Size([1, 2, 572, 572])

你一定会写UNet

  • 我的UNet的实现代码
  • 终于基于pytorch实现了UNet的模型。如果你还是不够清楚和明白UNet的实现过程，下面我分享几篇优秀的博文，你看完后一定会有深刻的理解，并会写自己的UNet模型。
    UNet的Pytorch实现 https://blog.csdn.net/kobayashi_/article/details/108951993
    Pytorch:Unet网络代码详解 https://cloud.tencent.com/developer/article/1633363
    UNet模型训练，深度解析 https://blog.csdn.net/weixin_43436958/article/details/107384695
    UNet 项目源码 https://github.com/milesial/Pytorch-UNet
@左左@右右
关注
  • 1
    点赞
  • 踩
  • 26
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
专栏目录
Python Unet网络结构pytorch简单实现+torchsummary可视化（可以直接运行）
12-21
Unet的网络结构： 根据该结构，用Pytorch实现Unet： # -*- coding: utf-8 -*- import torch import torch.nn as nn from torch.autograd import Variable import torchvision.datasets as dset import torchvision.transforms as transforms import torch.nn.functional as F import numpy as np import torch.utils.data as Data seed =
从零手写Unet代码及数据
smallworldxyl的博客
11-25 1910
数据集资源已上传，不需要积分。 一.网络模型（model.py） import torch import torch.nn as nn import torchvision.transforms.functional as TF class DoubleConv(nn.Module): def __init__(self,in_channels,out_channels): super(DoubleConv, self).__init__() self.conv=.
Unet环境配置pytroch框架
weixin_44751438的博客
02-21 5275
环境配置 先安装Anaconda，再安装cuda，配置torch环境 一、Anaconda安装 Anaconda的安装主要是为了方便环境管理，可以同时在一个电脑上安装多种环境，不同环境放置不同框架：pytorch、tensorflow、keras可以在不同的环境下安装，只需要使用conda create –n创建新环境即可。 1、Anaconda的下载 anaconda的下载： 如果想要安装最新的Anaconda，首先登录Anaconda的官网：https://www.anaconda.com
Skeletonization 基于UNet的骨架提取网络
m0_63642362的博客
02-25 1018
图像数据集由1725幅黑白图像组成，这些图像以256×256像素的便携式网络图形格式提供，分为1218幅训练图像、241幅验证图像和266幅测试图像。我们提供测试集和验证集中每个类的样本。有两种类型的图像：表示数据集中形状的形状图像和表示与形状图像对应的骨架的骨架图像。Pixel SkelNetOn 竞赛：像素骨架提取分析能对不同目标的二值化掩码进行细化，克服传统方法人工设置参数的不便。如何从少量样本中准确提取不同形状的的骨架，存在一定的挑战性。通过paddle提供的接口自定义数据加载方式。
语义分割系列6-Unet++（pytorch实现）
热门推荐
yumaomi的博客
05-17 2万+
本文介绍了Unet++网络，在pytorch框架上复现Unet++，并在Camvid数据集上进行训练。
语义分割系列2-Unet（pytorch实现）
yumaomi的博客
05-13 6973
本文详细介绍了Unet的网络结构和论文思路，在pytorch框架上复现了Unet结构，在Camvid数据集上进行测试。
深度学习：目标分割｜UNet网络模型及案例实现
m0_58475958的博客
08-15 1万+
1 UNet网络架构 UNet网络由左编码部分，右解码部分和下两个卷积+激活层组成 编码部分 从图中可知：架构中是由4个重复结构组成：2个3x3卷积层，非线形ReLU层和一个stride为2的2x2 max pooling层（图中的蓝箭头，红箭头） 每一次下采样特征通道的数量加倍 解码部分 和编码层类似，反卷积也有4个重复结构组成 每个重复结构前先使用反卷积，每次反卷积后特征通道数量减半，特征图大小加倍（绿箭头） 反卷积之后，反卷积的结果和编码部分对应步骤的特征图拼接起来（白/蓝块） 如果编
Pytorch QAT for UNet
牧羊女
06-20 1306
模型的QAT过程
pytorch快速入门与实战——三、Unet实现
无衣的博客
03-14 1万+
pytorch快速入门与实战——三、UNet适应输入AdUnet实现
基于Pytorch框架图像分割的红外弱小目标检测Python源码+数据集+模型+使用说明.zip
12-19
基于Pytorch框架图像分割的红外弱小目标检测Python源码+数据集+模型+使用说明.zip Training ```python python main.py unet train ``` Test ```python python main.py unet test image_dir ``` ### Evaluation ```...
基于PyTorch与PyTorch-Lightning进行人工神经网络模型与CNN模型的构建.zip
07-09
准备阶段，使用 PyTorch-Lightning 进行训练框架的搭建。 在{model_name}_main.py 入口脚本（例如 cnn_main.py）中设置 Global Seed 为 42，使用自定义的 MnistDataLoader 作为训练数据，使用 pl.Trainer()对模型...
UNet预训练模型（pytorch）
12-28
本预训练模型是基于pytorch框架，在医学图像上训练得到的。
基于Unet模型和并联注意力机制实现猫和老鼠动画片的语义分割
12-27
本项目使用爬虫技术获取数据集，经过数据清洗后使用百度EasyData平台实现数据标注，通过Pytorch框架搭建Unet模型，使用Resnet50作为编码器。使用迁移学习的思想，加载Resnet50在VOC数据集上训练好的预训练模型，从而...
ChangeDetection:使用PyTorch进行变更检测的框架
05-22
Change Detection一个用 PyTorch 编写的，专门针对变化检测 (Change Detection) 任务的模型框架。结果可视化（部分）Siamese_unet_conc + SzadaTODO参考写在前面为什么写这个项目？变化检测（Change Detection，CD）...
load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True
jacke121的专栏
04-28 284
load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True
支持向量机——pytorch与paddle实现支持向量机
Love you Joyfully的博客
04-28 1378
本文将深入探讨支持向量机的理论基础，并通过PyTorch和PaddlePaddle两个深度学习框架来展示如何实现支持向量机模型。我们将首先介绍支持向量机、优化的基本概念，这些数学工具是理解和实现支持向量机的基础。通过PyTorch和PaddlePaddle的代码示例，我们将展示如何设计、训练和评估一个支持向量机模型，从而让读者能够直观地理解并掌握这两种框架在机器学习问题中的应用。
深度学习框架pytorch：tensor.data和tensor.detach()的区别
huanfeng_AI的博客
04-24 227
现在来看大家还是推荐使用detach(）方法，然后当tensor的requires_grad=False情况下，tensor.data或者tensor.detach（）与tensor等价，而当tensor的requires_grad=True情况下，tensor.data为获取tensor中的数据。
Ubuntu如何更换 PyTorch 版本
最新发布
玩人工智能的辣条哥的博客
04-30 307
WLS2。
运行DeepSORT_YOLOv5_Pytorch时出现的问题
这个函数可导的博客
04-30 441
记录一下复现GitHub中以下这份代码时出现的一些问题。代码链接：https://github.com/HowieMa/DeepSORT_YOLOv5_Pytorch/tree/master。
pytorch框架UNet图像分割代码示例
04-04
以下是一个基于PyTorch框架实现的UNet图像分割代码示例： ```python import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): """Double convolution layer with ...

“相关推荐”对你有帮助么？

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值