Convert pixel wise class tensor to image segmentation
So as the label has format batch_size x 1 x height x width I can calculate cross entropy loss with: <SCODE>criterion = nn.CrossEntropyLoss()

# shape: batch_size x 1 x height x width = 1 x 1 x 256 x 256
inputs = autograd.Variable(images)
# shape: batch_size x 1 x height x width = 1 x 1 x 256 x 256
targets = autograd.Variable(labels)
# shape: batch_size x 1 x height x width = 1 x 256 x 256 x 256
outputs = model(inputs)

optimizer.zero_grad()
loss = criterion(outputs.view(-1, 256), targets.view(-1))
loss.backward()
optimizer.step()
<ECODE> I know that outputs[0][0][i][j| corresponds to the probability that the pixel at (i, j) belongs to class 1. So if want to transform outputs of shape 1 x 256 x 256 x 256 to 1 x 1 x 256 x 256 I would need to find the maximum (probability) of every pixel and assign it to the corresponding class value. I could do this manually by iterating over every class and pixel with numpy but I wonder if there is any better way using tensor operations?
bodokaiser
<SCODE>nllcrit = nn.NLLLoss2d(size_average=True) # need to write a functional interface for it
def criterion(input, target):
    return nllcrit(F.log_softmax(input), target)
<ECODE> <SCODE>output = output.permute(0, 2, 3, 1).view(-1, ncls).squeeze()
target = target.view(-1).squeeze()
<ECODE>
fmassa
bodokaiser
In torch, max and argmax are computed together and returned as a tuple. So outputs.max(0)[1] is the native torch way to do this.
jekbradbury
Any ideas what could be wrong? This happens with both models (UNet and 1-layer conv) so I guess there must be something wrong with loss or optimization.
bodokaiser
<SCODE>            # outputs.shape =(batch_size, n_classes, img_cols, img_rows) 
            outputs = outputs.permute(0, 2, 3, 1)
            # outputs.shape =(batch_size, img_cols, img_rows, n_classes) 
            outputs = outputs.resize(batch_size*img_cols*img_rows, n_classes)
            labels = labels.resize(batch_size*img_cols*img_rows)
            loss = F.cross_entropy(outputs, labels)
<ECODE>
meetshah1995
fmassa