三维unet用nn.CrossEntropyLoss做3分类

对于三维的 Unet 模型，如果要进行三分类，我们可以将输出层的通道数设置为 3，并使用 nn.CrossEntropyLoss 作为损失函数。

下面是一个示例代码：

import torch
import torch.nn as nn

class Unet3D(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Unet3D, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        
        # Encoder
        self.conv1 = nn.Conv3d(in_channels, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv3d(64, 128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv3d(128, 256, kernel_size=3, padding=1)
        self.conv4 = nn.Conv3d(256, 512, kernel_size=3, padding=1)
        
        # Decoder
        self.upconv1 = nn.ConvTranspose3d(512, 256, kernel_size=2, stride=2)
        self.conv5 = nn.Conv3d(512, 256, kernel_size=3, padding=1)
        self.upconv2 = nn.ConvTranspose3d(256, 128, kernel_size=2, stride=2)
        self.conv6 = nn.Conv3d(256, 128, kernel_size=3, padding=1)
        self.upconv3 = nn.ConvTranspose3d(128, 64, kernel_size=2, stride=2)
        self.conv7 = nn.Conv3d(128, 64, kernel_size=3, padding=1)
        
        # Output
        self.conv8 = nn.Conv3d(64, out_channels, kernel_size=1)
        
        # Activation functions
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        # Encoder
        x1 = self.relu(self.conv1(x))
        x2 = self.relu(self.conv2(x1))
        x3 = self.relu(self.conv3(x2))
        x4 = self.relu(self.conv4(x3))
        
        # Decoder
        x = self.upconv1(x4)
        x = torch.cat([x, x3], dim=1)
        x = self.relu(self.conv5(x))
        x = self.upconv2(x)
        x = torch.cat([x, x2], dim=1)
        x = self.relu(self.conv6(x))
        x = self.upconv3(x)
        x = torch.cat([x, x1], dim=1)
        x = self.relu(self.conv7(x))
        
        # Output
        x = self.conv8(x)
        x = self.softmax(x)
        
        return x

# 定义模型
model = Unet3D(in_channels=1, out_channels=3)

# 定义损失函数
loss_fn = nn.CrossEntropyLoss()

# 定义优化器
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

在这个例子中，我们定义了一个包含两个卷积层和三个反卷积层的 Unet3D 模型，并将输出通道数设置为 3。我们使用 nn.CrossEntropyLoss 作为损失函数，并将模型的输出结果和真实的标签传递给损失函数，计算损失并反向传播更新参数。

需要注意的是，nn.CrossEntropyLoss 要求模型的最后一层输出的大小必须等于类别的数量。在本例中，我们有三个类别，因此最后一层的输出大小为 3。同时，由于我们的模型输出的是一个概率分布，因此我们还需要使用 softmax 函数将输出转换为概率。