Example CrossEntropyLoss for 3D semantic segmentation in pytorch
I have a network performing 3D convolutions on a 5D input tensor. The output of my network if of size (1, 12, 60, 36, 60) corresponding to ( BatchSize, NumClasses, x-dim, y-dim, z-dim). I need to compute a voxel-wise cross entropy loss. However I keep on getting errors.

When trying to compute cross entropy loss using torch.nn.CrossEntropyLoss(), I keep on getting the following error message:  

RuntimeError: multi-target not supported at .../src/THCUNN/generic/ClassNLLCriterion.cu:16


here is the extract of my code:  

import torch
import torch.nn as nn
from torch.autograd import Variable
criterion = torch.nn.CrossEntropyLoss()
images = Variable(torch.randn(1, 12, 60, 36, 60)).cuda()
labels = Variable(torch.zeros(1, 12, 60, 36, 60).random_(2)).long().cuda()
loss = criterion(images.view(1,-1), labels.view(1,-1))


Same happens when I create a one-hot tensor for the labels:  

nclasses = 12
labels = (np.random.randint(0,12,(1,60,36,60))) # Random labels with values between [0..11]
labels = (np.arange(nclasses) == labels[..., None] - 1).astype(int) # Converts labels to one_hot_tensor
a = np.transpose(labels,(0,4,3,2,1)) #  Reorder dimensions to match shape of "images" ([1, 12, 60, 36, 60])
b = Variable(torch.from_numpy(a)).cuda()
loss = criterion(images.view(1,-1), b.view(1,-1))


Any idea what I'm doing wrong?
Can someone provide an example of computing cross entropy on a 5D output tensor?

Just checked some implementation (fcn) for 2D semantic segmentation, and tried to adapt it to 3D semantic segmentation. No guarantee that this is correct, I'll have to double check...  

 import torch
 import torch.nn.functional as F
 def cross_entropy3d(input, target, weight=None, size_average=True):
    # input: (n, c, h, w, z), target: (n, h, w, z)
    n, c, h, w , z = input.size()
    # log_p: (n, c, h, w, z)
    log_p = F.log_softmax(input, dim=1)
    # log_p: (n*h*w*z, c)
    log_p = log_p.permute(0, 4, 3, 2, 1).contiguous().view(-1, c) # make class dimension last dimension
    log_p = log_p[target.view(n, h, w, z, 1).repeat(1, 1, 1, 1, c) &gt;= 0] # this looks wrong -&gt; Should rather be a one-hot vector
    log_p = log_p.view(-1, c)
    # target: (n*h*w*z,)
    mask = target &gt;= 0
    target = target[mask]
    loss = F.nll_loss(log_p, target.view(-1), weight=weight, size_average=False)
    if size_average:
        loss /= mask.data.sum()
    return loss
images = Variable(torch.randn(5, 3, 16, 16, 16))
labels = Variable(torch.LongTensor(5, 16, 16, 16).random_(3))
cross_entropy3d(images, labels, weight=None, size_average=True)

