Unet代码实现(PyTorch)

張張-張張

已于 2023-08-01 16:38:46 修改

阅读量2.8k

点赞数 1

分类专栏：语义分割文章标签：深度学习 UNet 语义分割

于 2023-04-22 14:31:23 首次发布

本文链接：https://blog.csdn.net/proceduralman/article/details/130304557

版权

该代码实现了一个基于PyTorch的UNet模型，用于图像分割任务。模型包括下采样和上采样模块，每个模块由卷积层、批量归一化、LeakyReLU激活函数和Dropout层组成。在上采样过程中，采用了双线性插值和特征拼接来恢复分辨率。

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

import torch
import torch.nn as nn
from torch.nn import functional as F


# 基本卷积块， 长宽不变，in_channels -> out_channels
class Conv(nn.Module):
    def __init__(self, C_in, C_out):
        super(Conv, self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(C_in, C_out, 3, 1, 1),
            nn.BatchNorm2d(C_out),
            # 防止过拟合
            nn.Dropout(0.3),
            nn.LeakyReLU(),

            nn.Conv2d(C_out, C_out, 3, 1, 1),
            nn.BatchNorm2d(C_out),
            # 防止过拟合
            nn.Dropout(0.4),
            nn.LeakyReLU(),
        )

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

# 下采样模块，长宽下采样2倍，通道数不变
class DownSampling(nn.Module):
    def __init__(self, C):
        super(DownSampling, self).__init__()
        self.Down = nn.Sequential(
            # 使用卷积进行2倍的下采样，通道数不变
            nn.Conv2d(C, C, 3, 2, 1),
            nn.LeakyReLU()
        )

    def forward(self, x):
        return s