During the process of attempting to extract table lines (which I divided into horizontal and vertical lines), I found that when I used CrossEntropyLoss, the loss value remained at a constant level without changing. However, when I used BCELoss, the loss gradually decreased with each iteration.
For example:
Input: (batch_size, 3, 1024, 1024)
Output: (batch_size, 2, 1024, 1024)
Target: (batch_size, 2, 1024, 1024)
nn.CrossEntropyLoss can be used for multi-class classification use cases (where each sample belongs to a single class) while nn.BCEWithLogitsLoss (don’t use nn.BCELoss as it has a lower numerical stability compared to nn.BCEWithLogitsLoss) can be used for multi-label classification use cases (where each sample belongs to 0, 1, or multiple classes).
I don’t know which use case you are working on.
For example, in an image, I would first mark the lines of the table (that is, all parts that belong to the table are marked as 1, and those that do not belong to the table are marked as 0). Then I use a model to determine whether each pixel belongs to a line.
In this case it sounds like a multi-class segmentation use case and nn.CrossEntropyLoss would be the commonly used criterion.
The target would then contain class indices in [0, 1] and have the shape [batch_size, height, width] (note the missing class dimension) or would contain “soft” targets (i.e. floating point values between 0 - 1) in the same shape as the model output.
1 Like