tensorflow - softmax_cross_entropy_with_logits に対する PyTorch の等価性

Question

TensorFlow の同等の PyTorch 損失関数があるのだろうかと思っていましたsoftmax_cross_entropy_with_logits。

Answer 1

いくつかのスレッドでポインターをたどると、最終的に次の変換になりました。他の誰かがこのスレッドに陥った場合に備えて、ここに私の解決策を投稿します。ここから変更され、このコンテキスト内で期待どおりに動作します。

# pred is the prediction with shape [C, H*W]
# gt is the target with shape [H*W]
# idx is the boolean array on H*W for masking

# Tensorflow version
loss = tf.nn.sparse_softmax_cross_entropy_with_logits( \
          logits=tf.boolean_mask(pred, idx), \
          labels=tf.boolean_mask(gt, idx)))

# Pytorch version       
logp = torch.nn.functional.log_softmax(pred[idx])
logpy = torch.gather(logp, 1, Variable(gt[idx].view(-1,1)))
loss = -(logpy).mean()