In training a fully convolutional network for semantic segmentation, the loss function is fundamentally similar to that used in standard image classification tasks. However, because the network predicts a distinct class for every individual pixel, these predictions are represented by the output channels of the final transposed convolutional layer. Consequently, when calculating the loss—typically a cross-entropy loss—the computation must explicitly specify the channel dimension to correctly evaluate the predicted probabilities against the ground-truth classes for each spatial position.

# PyTorch
def loss(inputs, targets):
    return F.cross_entropy(inputs, targets, reduction='none').mean(1).mean(1)

# MXNet
loss = gluon.loss.SoftmaxCrossEntropyLoss(axis=1)