Unet多类别分割的输出

下面是一个使用PyTorch实现Unet多类别分割的输出的示例代码：

import torch
import torch.nn as nn

class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()
        
        # 编码器
        self.enc1 = self.encoder(3, 64)
        self.enc2 = self.encoder(64, 128)
        self.enc3 = self.encoder(128, 256)
        self.enc4 = self.encoder(256, 512)
        
        # 解码器
        self.dec1 = self.decoder(512, 256)
        self.dec2 = self.decoder(256, 128)
        self.dec3 = self.decoder(128, 64)
        
        # 分类器
        self.classifier = nn.Conv2d(64, num_classes, kernel_size=1)
        
    def encoder(self, in_channels, out_channels):
        encoder = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True)
        )
        return encoder
    
    def decoder(self, in_channels, out_channels):
        decoder = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(out_channels, in_channels, kernel_size=2, stride=2)
        )
        return decoder
    
    def forward(self, x):
        # 编码器
        enc1 = self.enc1(x)
        enc2 = self.enc2(enc1)
        enc3 = self.enc3(enc2)
        enc4 = self.enc4(enc3)
        
        # 解码器
        dec1 = self.dec1(enc4)
        dec2 = self.dec2(torch.cat([dec1, enc3], dim=1))
        dec3 = self.dec3(torch.cat([dec2, enc2], dim=1))
        
        # 分类器
        output = self.classifier(dec3)
        return output

# 示例使用7个类别
num_classes = 7

# 创建模型实例
model = UNet()

# 创建输入数据（假设输入尺寸为[batch_size, 3, height, width]）
batch_size = 1
height = 256
width = 256
input_data = torch.randn(batch_size, 3, height, width)

# 前向传播
output = model(input_data)

# 输出的尺寸为[batch_size, num_classes, height, width]
print(output.shape)

这是一个简化的Unet模型，包括编码器、解码器和分类器部分。编码器由四个卷积层组成，解码器由三个卷积层和一个反卷积层组成。分类器部分使用1x1卷积层将解码器的输出映射到指定类别数。在前向传播过程中，输入数据通过编码器和解码器进行处理，最后通过分类器进行分类，输出为每个像素点属于每个类别的概率。输出的尺寸为[batch_size, num_classes, height, width]，其中num_classes为类别数。