Hi Stark!



 stark:

What’re the options if one wants to use Dice or IOU as a loss function?


Let me imagine that your use case is something like multi-class
semantic segmentation.  Your target is an “image” of class labels,
that is this pixel is “class-3,” that pixel is “class-1,” and so on.  The
prediction from your network is a set of nClass raw-score logits
for each pixel, with shape [nBatch, nClass, height, width].
Softmax (perhaps better named SoftOneHotEncodedArgMax)
will differentiably push your predictions towards 0.0 and 1.0.
So use soft_labels = torch.softmax (alpha * pred, dim = 1)
alpha can be used to tune how “hard” your soft one-hot labels
become – as you increase alpha you will push the labels closer
to being exactly zero or one.  But don’t make alpha too large – if
the labels saturate and become exactly zero and one you will lose
useful differentiability, and as they get close to saturating, your
differentiability will become less and less useful.
Probably the most straightforward way to calculate your “soft”
intersection-over-union will be to one-hot your integer-class-label
target so that it also has shape [nBatch, nClass, height, width].
Letting pred be your soft-label predictions and targ be your
one-hotted class labels, then:
iou = (pred * targ).sum (dim = (2, 3)) / (pred + targ).sum (dim = (2, 3))

(If pred and targ were boolean – that is just 0 or 1 – pred * targ
would be the boolean intersection and pred + targ the union.)
iou will now have shape [nBatch, nClass] and each element of it
will be the soft IOU for a specific class and a specific sample in the
batch.
Best.
K. Frank