How to implement cross entropy for binary segmentation using symbol only?

[John1231983]

Given a softmax classified feature ff (Bx2xHxW). And a target label size of Bx1xHxW . I want to implement cross entropy loss using symbol only. This is my implementation

# target size of Bx1xHxW
target_squeeze = mx.symbol.squeeze(target, axis=1) #size of BxHxW
target_squeeze = mx.sym.one_hot(target_squeeze, depth = 2, on_value = -1.0, off_value = 0.0) 
# Transpose from BxHxWx2 to Bx2xHxW
 target_squeeze = mx.symbol.transpose(target_squeeze, axes=(0,3,1,2))
# Get log of feature f
f_log  = mx.sym.log(f)
batch_size =32
f_sum = mx.symbol.sum(target_squeeze * f_log)/batch_size
f_sum = mx.symbol.MakeLoss(f_sum, name = 'loss_ce')

Is my implementation correct? If not, please correct it for me. Thanks in advantage

[John1231983]

Any suggestion?

[wkcn]

Hi @John1231983 , using mx.nd.pick will simplify the code.

import mxnet as mx

B, C, H, W = 2, 2, 3, 4

x = mx.random.uniform(-1, 1, shape=(B, C, H, W)) 
target = mx.random.randint(0, C, shape=(B, 1, H, W)) 

f = mx.nd.softmax(x)
# target size of Bx1xHxW
target_squeeze = mx.nd.squeeze(target, axis=1) #size of BxHxW
target_squeeze = mx.nd.one_hot(target_squeeze, depth = 2, on_value = -1.0, off_value = 0.0) 
# Transpose from BxHxWx2 to Bx2xHxW
target_squeeze = mx.nd.transpose(target_squeeze, axes=(0,3,1,2))
# Get log of feature f
f_log  = mx.nd.log(f)
batch_size =32
f_sum = mx.nd.sum(target_squeeze * f_log)/batch_size
print(f_sum)

lscore = -mx.nd.log_softmax(x)
target_squeeze = mx.nd.squeeze(target, axis=1) #size of BxHxW
t_sum = mx.nd.pick(lscore, target_squeeze, axis=1).sum() / batch_size
print(t_sum) # t_sum == f_sum

[John1231983]

@wkcn thanks so much for your answer. Could you please verify my implementation. Does it correct, should I use axis=1 in sum and then use .mean()

[wkcn]

@John1231983
Your implementation is correct.
It is available to use data.sum(axis=1).mean(), which is equal to data.sum() / len(data).