unet代码详解pytorch

UNet是一种用于图像分割的卷积神经网络，其结构类似于自编码器。在这个网络中，输入图像被编码成一个较小的向量，然后被解码为分割图像。UNet的特点是具有对称的编解码结构，同时也具备跳跃连接，这使得网络能够更好地捕捉到不同层次的特征信息。

下面是一个基于PyTorch实现的UNet模型的代码详解：

首先，我们需要导入所需的库：

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

接下来，我们定义UNet的编码器和解码器部分：

class EncoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(EncoderBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        out = self.pool(x)
        return out, x

class DecoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DecoderBlock, self).__init__()
        self.upconv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channels)

    def forward(self, x, encoder_output):
        x = self.upconv(x)
        x = torch.cat([x, encoder_output], dim=1)
        x = F.relu(self.bn1(self.conv1(x)))
        out = F.relu(self.bn2(self.conv2(x)))
        return out

EncoderBlock和DecoderBlock都是基本的卷积块，其中EncoderBlock用于向下采样图像，DecoderBlock用于向上采样图像。在EncoderBlock中，我们使用2个卷积层和1个最大池化层来减小图像的大小。在DecoderBlock中，我们使用一个转置卷积层和2个卷积层来增加图像的大小。在解码器中，我们还使用了跳跃连接，将编码器输出的特征图与解码器的输入特征图连接起来。

接下来，我们定义完整的UNet模型：

class UNet(nn.Module):
    def __init__(self, num_classes=1):
        super(UNet, self).__init__()
        self.enc1 = EncoderBlock(3, 64)
        self.enc2 = EncoderBlock(64, 128)
        self.enc3 = EncoderBlock(128, 256)
        self.enc4 = EncoderBlock(256, 512)

        self.center = nn.Conv2d(512, 1024, kernel_size=3, padding=1)
        self.bn = nn.BatchNorm2d(1024)

        self.dec4 = DecoderBlock(1024, 512)
        self.dec3 = DecoderBlock(512, 256)
        self.dec2 = DecoderBlock(256, 128)
        self.dec1 = DecoderBlock(128, 64)

        self.final = nn.Conv2d(64, num_classes, kernel_size=1)

    def forward(self, x):
        enc1, out1 = self.enc1(x)
        enc2, out2 = self.enc2(enc1)
        enc3, out3 = self.enc3(enc2)
        enc4, out4 = self.enc4(enc3)

        center = self.center(enc4)
        center = self.bn(center)

        dec4 = self.dec4(center, out4)
        dec3 = self.dec3(dec4, out3)
        dec2 = self.dec2(dec3, out2)
        dec1 = self.dec1(dec2, out1)

        final = self.final(dec1)
        return final

其中，我们将4个EncoderBlock和4个DecoderBlock连接起来，中间加入了一个卷积层和BatchNormalization层。

最后，我们可以使用该模型进行图像分割：

model = UNet(num_classes=2)
inputs = torch.randn((1, 3, 256, 256))
outputs = model(inputs)
print(outputs.shape)

这里我们使用了一个大小为256x256的RGB图像进行测试，并输出了模型的输出形状。