如何为像素分类做软最大值

Question

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

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, 2)
return x

在最后一层中，我设计了 conv13，使其为每个类生成 2 个通道。

由于我使用的是 softmax，因此我期望 2 个独立通道上相同索引的值总和等于 1。

例如，假设输出图像是(2{通道}，4,4(。所以我期待

图像[ 通道 1 ][0

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

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

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

for batch, data in enumerate(trainloader, 0):
inputs , labels = data
inputs, labels = inputs.to(device), labels.to(device)

optimizer.zero_grad()
outputs = net(inputs)
loss = rmse(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()

predicted = outputs.data
predicted = predicted.to('cpu')
predicted_img = predicted.numpy()
predicted_img = np.reshape(predicted_img,(2, 4, 4))
print(predicted_img[0])
print(predicted_img[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( ！= 1

我该如何解决？

Answer 1

请检查我代码的最后一行..基本上你的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)

希望有帮助。