Computer vision 如何使用softmax进行像素分类

我的目标是使用像素分类进行灰度图像分割。所以我有两个标签0和1。我在pytorch中创建了一个网络，如下所示

类网络（nn.Module）：
定义初始化（自）：
超级（网络，自我）。\uuuu初始化
self.up=nn.Upsample（比例系数=2，模式='nearest'）
self.conv11=nn.Conv2d（1128，内核大小=3，填充=1）
self.conv12=nn.Conv2d（128256，内核大小=3，填充=1）
self.conv13=nn.Conv2d（256，2，内核大小=3，填充=1）
def前进（自身，x）：
in_size=x.size（0）
x=F.relu（自转换11（x））
x=F.relu（自转换12（x））
x=F.relu（自转换13（x））
x=F.softmax（x，2）
返回x

在最后一层中，我设计了conv13 In，以便为每个类生成2个通道

因为我使用的是softmax，所以我希望两个单独通道上相同索引的值之和等于1

例如，假设输出图像是（2{channel}，4，4）。所以我期待着

图像[频道1][0][0]+图像[频道2][0][0]=1

但我得到的输出是0.0015，甚至不接近1。我如何使用softmax来预测频道

为了检查这一点，我使用了以下代码

对于批次，枚举中的数据（火车装载机，0）：
输入、标签=数据
输入，标签=输入。到（设备），标签。到（设备）
optimizer.zero_grad（）
输出=净（输入）
损耗=rmse（输出、标签）
loss.backward（）
optimizer.step（）
运行损失+=损失。项目（）
预测=输出。数据
预测=预测。到（'cpu'）
predicted\u img=predicted.numpy（）
预测的\u img=np.重塑（预测的\u img，（2,4,4））
打印（预测的\u img[0]）
打印（预测的图像[1]）

那些指纹显示了这一点

[[**0.2762002**0.13305853 0.2510342 0.23114938]
[0.26812425 0.28500515 0.05682982 0.15851443]
[0.1640967  0.5409352  0.43547812 0.44782472]
[0.29157883 0.0410011 0.2566578 0.16251141]

[[**0.23052207**0.868455 0.43436486 0.0684725]
[0.18001427 0.02341573 0.0727293  0.2525512 ]
[0.06587404 0.04974682 0.3773188  0.6559266 ]
[0.5235896 0.05838248 0.11558701 0.02304965]]

很明显，相应的元素的总和不是1

0.2762002（指数0,0）+0.23052207（指数0,0）！=一,

---

如何修复它？

请检查我代码的最后一行。。基本上，softmax的尺寸是错误的

class Net(nn.Module):

  def __init__(self):
        super(Net, self).__init__()

        self.up = nn.Upsample(scale_factor=2, mode='nearest')

        self.conv11 = nn.Conv2d(1, 128, kernel_size=3, padding=1)
        self.conv12 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.conv13 = nn.Conv2d(256, 2, kernel_size=3, padding=1)  



    def forward(self, x):
        in_size = x.size(0)

        x = F.relu(self.conv11(x))
        x = F.relu(self.conv12(x))
        x = F.relu(self.conv13(x))

        x = F.softmax(x, 1) #this line is changed

        return x

net = Net()
inputs = torch.rand(1,1,4,4)
out = net (Variable(inputs))
print (out)
out.sum(dim=1)

---

希望有帮助

为什么tensorflow和pytorch标签都有？编辑！似乎stackoverflow自动建议了tensorflow标签，我提前提交了问题，但没有删除它！非常感谢你。基本上我刚才就知道了。我给F.softmax（x，2）2上了2节课。但后来意识到，参数询问在哪个维度上执行softmax，它不询问类的数量。因为我希望它跨越通道，所以值应该是1，因为火炬张量的形状像[批次、通道、高度、宽度]。