U-Net的pytorch实现

最新推荐文章于 2024-09-15 07:07:27 发布

weixin_39266208

最新推荐文章于 2024-09-15 07:07:27 发布

阅读量3k

点赞数 1

分类专栏： pytorch 文章标签：计算机视觉 pytorch

本文链接：https://blog.csdn.net/weixin_39266208/article/details/121366475

版权

本文详细介绍了U-Net网络结构，包括其在细胞识别任务中的应用和作者的PyTorch实现。文章阐述了网络的对称结构，下采样和上采样的过程，并讨论了如何处理输入和输出尺寸的问题。此外，还提到了训练时的数据处理技巧，如输入图像的大小选择和边缘处理，以及验证阶段的注意事项。尽管没有提供严格的验证结果，但文章为理解U-Net网络提供了一定的参考。

摘要由CSDN通过智能技术生成

简介

查看U-Net论文请点击此处。U-Net最初是用于细胞识别的，针对少量的训练数据，作者通过数据增强等方式，实现了很好的效果，获得了比赛冠军，U-Net经过修改也可以用于其他用途。
U-Net论文作者提供了caffe的版本，github上也已经有人提供了pytorch的版本，但是经过了修改，本文提供的实现忠于论文的描述，没有研究过作者提供的版本，所以不保证本实现和作者的实现足够接近，仅是个人对论文的理解，如有错误，还望指正，欢迎留言，另外，由于没有对应的训练数据，所以没有严格的验证。

正文

代码如下，代码后面有说明

import torch
import torch.nn as nn
from torchsummary import summary

class MyUnetDown(nn.Module):
    def __init__(self, in_channels):
        super(MyUnetDown, self).__init__()
        
        self.down_and_conv = nn.Sequential(
            nn.MaxPool2d(
                kernel_size = 2,
                stride=2,
            ),
            nn.Conv2d(
                in_channels = in_channels,
                out_channels = in_channels*2,
                kernel_size = 3,
                stride=1,
                padding=0, # 论文中说unpadded convolutions
                bias=False,
            ),
            nn.ReLU(),
            nn.Conv2d(
                in_channels = in_channels*2,
                out_channels = in_channels*2,
                kernel_size = 3,
                stride=1,
                padding=0,
                bias=False,
            ),
            nn.ReLU(),
        ) # 这个地方如果不小心写了个逗号，就玩完了!!!
        
    def forward(self, x):
        return self.down_and_conv(x)
    

class MyUnetUp(nn.Module):
    def __init__(self, in_channels, cropsize):
        super(MyUnetUp, self).__init__()
        
        # 和最大池化对应，如果遇到奇数，没有padding最大池化会丢弃多余的，所以这里out_padding=0
        # 最大池化大小(形状)上相当于卷积kernel=2, stride=2, 这里是大小(形状)上的逆过程
        # 参数和对应的Conv2d相同， stride=2 卷积回去有歧义， out_padding区分歧义，这里设置为0,对应偶数
        # 详情参见ConvTranspose2d文档以及文档中给的ConvTranspose2d的图示链接
        #
        # ConvTranspose2d根据参数 dilation * (kernel_size - 1) - padding计