CrossEntropyLoss for 3D and higher

[wadimkehl]

Going hand-in-hand with #1020 and #1260 it would be nice to have the same multi-dim behavior for the CrossEntropyLoss as is now for softmax. Similar to https://www.tensorflow.org/api_docs/python/tf/nn/sparse_softmax_cross_entropy_with_logits
where the loss is is evaluated along the last provided dimension, i.e.

labels: shape [d_0, d_1, ..., d_{r-1}]
logits: shape [d_0, d_1, ..., d_{r-1}, num_classes]

At the moment one has to reshape into 2D before and back again after the operation.
It would not require a signature change and would not break current code that use 2D input into the loss.
Any thoughts?

[soumith]

agreed, would be good to have. We'll get this done.

[zou3519]

Looking into this

[soumith]

fixed via #4035

[John1231983]

@soumith and @zou3519 : Thanks for it, I have same problem. I am implementing 3D Unet loss function.

• The labels size is batch_size x depth x height x width
• The logits size is batch_size x number_class x depth x height x width

How can I use CrossEntropyLoss to compute the loss? this is my current way, but not sure the correction

criterion = nn.CrossEntropyLoss()
for epoch in range (num_epoch): 
     loss = 0
     for deep_slice in range(depth):
            loss += criterion(logits[:, :, deep_slice], labels[:, deep_slice])
     loss=loss_seg/depth
     loss.backward()

Another way may be converted to 2D

B,C,D,H,W=outputs.size()
logits_2d=logits.view(B,C,D, -1)
labels_2d = labels.view(B, D,-1)
loss = criterion(logits_2d, labels_2d)

Thanks!

[zou3519]

@John1231983 a better place to ask questions is the forums; your question will have higher visibility there than on a closed issue.

Does criterion(logits, labels) not work for you?

[John1231983]

@zou3519; yes, it does not work because CrossEntropyLoss only work when my input size like B x C x H x W but my case is Bx C x D x Hx W
Update: sorry it works now with my input size. I think I have upgrade pytorch