UNET网络双解码器模型代码

以下是基于PyTorch的UNET网络双解码器模型代码：

import torch
import torch.nn as nn

class DoubleDecoderUNet(nn.Module):
    def __init__(self):
        super(DoubleDecoderUNet, self).__init__()

        # Encoder
        self.encoder_conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.encoder_bn1 = nn.BatchNorm2d(64)
        self.encoder_relu1 = nn.ReLU(inplace=True)
        self.encoder_maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.encoder_conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.encoder_bn2 = nn.BatchNorm2d(128)
        self.encoder_relu2 = nn.ReLU(inplace=True)
        self.encoder_maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.encoder_conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.encoder_bn3 = nn.BatchNorm2d(256)
        self.encoder_relu3 = nn.ReLU(inplace=True)
        self.encoder_maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.encoder_conv4 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
        self.encoder_bn4 = nn.BatchNorm2d(512)
        self.encoder_relu4 = nn.ReLU(inplace=True)
        self.encoder_maxpool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.encoder_conv5 = nn.Conv2d(512, 1024, kernel_size=3, padding=1)
        self.encoder_bn5 = nn.BatchNorm2d(1024)
        self.encoder_relu5 = nn.ReLU(inplace=True)

        # Decoder 1
        self.decoder1_upsampling = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
        self.decoder1_conv1 = nn.Conv2d(1024, 512, kernel_size=3, padding=1)
        self.decoder1_bn1 = nn.BatchNorm2d(512)
        self.decoder1_relu1 = nn.ReLU(inplace=True)
        self.decoder1_conv2 = nn.Conv2d(512, 512, kernel_size=3, padding=1)
        self.decoder1_bn2 = nn.BatchNorm2d(512)
        self.decoder1_relu2 = nn.ReLU(inplace=True)

        # Decoder 2
        self.decoder2_upsampling = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.decoder2_conv1 = nn.Conv2d(512, 256, kernel_size=3, padding=1)
        self.decoder2_bn1 = nn.BatchNorm2d(256)
        self.decoder2_relu1 = nn.ReLU(inplace=True)
        self.decoder2_conv2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.decoder2_bn2 = nn.BatchNorm2d(256)
        self.decoder2_relu2 = nn.ReLU(inplace=True)

        # Final convolution
        self.final_conv = nn.Conv2d(256, 1, kernel_size=1)

    def forward(self, x):
        # Encoder
        x = self.encoder_conv1(x)
        x = self.encoder_bn1(x)
        x = self.encoder_relu1(x)
        x1 = x.clone()
        x = self.encoder_maxpool1(x)

        x = self.encoder_conv2(x)
        x = self.encoder_bn2(x)
        x = self.encoder_relu2(x)
        x2 = x.clone()
        x = self.encoder_maxpool2(x)

        x = self.encoder_conv3(x)
        x = self.encoder_bn3(x)
        x = self.encoder_relu3(x)
        x3 = x.clone()
        x = self.encoder_maxpool3(x)

        x = self.encoder_conv4(x)
        x = self.encoder_bn4(x)
        x = self.encoder_relu4(x)
        x4 = x.clone()
        x = self.encoder_maxpool4(x)

        x = self.encoder_conv5(x)
        x = self.encoder_bn5(x)
        x = self.encoder_relu5(x)

        # Decoder 1
        x = self.decoder1_upsampling(x)
        x = torch.cat([x, x4], dim=1)

        x = self.decoder1_conv1(x)
        x = self.decoder1_bn1(x)
        x = self.decoder1_relu1(x)

        x = self.decoder1_conv2(x)
        x = self.decoder1_bn2(x)
        x = self.decoder1_relu2(x)

        # Decoder 2
        x = self.decoder2_upsampling(x)
        x = torch.cat([x, x3], dim=1)

        x = self.decoder2_conv1(x)
        x = self.decoder2_bn1(x)
        x = self.decoder2_relu1(x)

        x = self.decoder2_conv2(x)
        x = self.decoder2_bn2(x)
        x = self.decoder2_relu2(x)

        # Final convolution
        x = self.final_conv(x)

        return x

这个模型有一个双解码器结构，其中第一个解码器的输出被用作第二个解码器的输入。这种结构可以更好地捕捉不同尺度的特征。