PyTorch中如何将(24,2,224,224)张量转为二值分割结果张量？

刚好做过类似的二值分割任务，你的需求可以通过PyTorch内置的torch.argmax()函数快速实现，具体步骤和代码如下：

核心思路 #

你的输入张量维度是(24, 2, 224, 224)，其中第1维（索引从0开始）对应前景/背景的概率。我们需要在每个像素位置比较这两个概率值，取概率更高的类别索引（0或1），然后调整维度到你需要的(24, 1, 224, 224)。

具体代码实现 #

import torch

# 模拟你的CNN输出张量（形状24,2,224,224，已做softmax确保概率和为1）
output_probs = torch.randn(24, 2, 224, 224).softmax(dim=1)

# 1. 取每个位置概率最大的类别索引，此时形状变为(24, 224, 224)
mask_indices = output_probs.argmax(dim=1)
# 2. 在第1维插入一个维度，得到目标形状(24, 1, 224, 224)
binary_mask = mask_indices.unsqueeze(1)

# 验证形状
print(binary_mask.shape)  # 输出: torch.Size([24, 1, 224, 224])

细节说明 #

• argmax(dim=1)的作用：指定在第1维（类别维度）上取最大值的索引，这样每个像素位置会得到0（背景概率更高）或1（前景概率更高）的结果。
• unsqueeze(1)的作用：argmax会把类别维度去掉，所以需要通过这个函数在原类别维度的位置插入一个新维度，恢复到你需要的张量结构。
• 数据类型调整：如果需要让输出掩码的类型和原概率张量一致（比如浮点型），可以在最后加上.float()：

  binary_mask = mask_indices.unsqueeze(1).float()
  

• 优化提示：如果你的CNN输出是未经过softmax的logits，其实可以直接对logits使用argmax——因为softmax是单调递增函数，最大值的位置不会改变，这样能节省一步计算。

内容的提问来源于stack exchange，提问作者Jimbo