Pytorch入门——用UNet网络做图像分割

最新推荐文章于 2025-06-24 09:28:37 发布

原创

最新推荐文章于 2025-06-24 09:28:37 发布 · 3.1w 阅读

354 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习 #神经网络

这篇博客介绍如何使用PyTorch从头实现UNet网络进行图像分割，通过阅读相关paper并参考其他博主的代码，作者逐步解析了网络结构和运行流程，同时提供了数据集链接和关键代码段。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

最近看的paper里的pytorch代码太复杂，我之前也没接触过pytorch，遂决定先自己实现一个基础的裸代码，这样走一遍，对跑网络的基本流程和一些常用的基础函数的印象会更深刻。

本文的代码和数据主要来自pytorch笔记：05)UNet网络简单实现_Javis486的专栏-优快云博客，

附上该博主的github地址：https://github.com/JavisPeng/u_net_liver

并在自己的理解的基础上做了一些改动，以及加了大量注释。

如有错误，欢迎指出。

unet.py（实现unet网络）

import torch.nn as nn
import torch

class DoubleConv(nn.Module):
    def __init__(self,in_ch,out_ch):
        super(DoubleConv,self).__init__()
        self.conv = nn.Sequential(
                nn.Conv2d(in_ch,out_ch,3,padding=1),#in_ch、out_ch是通道数
                nn.BatchNorm2d(out_ch),
                nn.ReLU(inplace = True),
                nn.Conv2d(out_ch,out_ch,3,padding=1),
                nn.BatchNorm2d(out_ch),
                nn.ReLU(inplace = True)  
            )
    def forward(self,x):
        return self.conv(x)


class UNet(nn.Module):
    def __init__(self,in_ch,out_ch):
        super(UNet,self).__init__()
        self.conv1 = DoubleConv(in_ch,64)
        self.pool1 = nn.MaxPool2d(2)#每次把图像尺寸缩小一半
        self.conv2 = DoubleConv(64,128)
        self.pool2 = nn.MaxPool2d(2)
        self.conv3 = DoubleConv(128,256)
        self.pool3 = nn.MaxPool2d(2)
        self.conv4 = DoubleConv(256,512)
        self.pool4 = nn.MaxPool2d(2)
        self.c