将 UNet 模型输出从 logits 转换为分割掩码

Question

我最近一直在从事某些 CT 扫描的分割任务。我决定使用

Python

作为语言，

ResidualUNet

模型架构在 MONAI 中实现，由 DiceLoss 评估。 一切都很顺利，直到训练有素的模型的干扰。这个 ResidualUNet 没有 Softmax 或 Sigmoid 层，因此它输出原始 logits，是从

-30

到

12

左右的浮点数。 如何正确转换这些逻辑以获得属于每个类别的每个像素的概率？ 输入用

BCHWD

表示，是

1x1x256x256x256

，当然是图像本身，而输出是

1x9x256x256x256

，每个通道是不同类别的掩码。 推理代码看起来不太像这样：

device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)

model = Unet(
    spatial_dims = 3,
    in_channels = 1,
    out_channels = 9,
    channels = (8, 16, 32, 64, 128, 256),
    strides = (2,2,2,2,2),
    num_res_units = 4,
    norm = Norm.INSTANCE
).to(device)

model.load_state_dict(torch.load(PATH_TO_SAVED_MODEL_OBJECT)["model_state_dict"])

inputs, labels = next(iter(validation_dataloader))   # obtaining only one image - batch_size=1
inferer = SimpleInferer()   # simple inferer from monai.inferers
inputs = inputs.to(device)  # pass to GPU
labels = labels.to(device)  # pass to GPU
pred = inferer(inputs = inputs, network=model)
pred = pred.detach().cpu().numpy() # conversion to numpy array for viewing purposes

预先感谢您的协助。

Answer 1

如果您需要多类分类，以便每个像素可以属于多个类，则可以使用 sigmoid，即

prob = torch.sigmoid(pred)

，它将为您提供每个像素独立属于每个类的概率，如果您需要单类分类，您可以使用 softmax，即

prob = torch.softmax(pred, dim=1)

。这两者都会将 logits 转换为有效的概率分布。