Apply cross entropy loss with custom weight map

username: Rishav_Sapahia
cooked: <p>I am solving multi-class segmentation problem using u-net architecture. As specified in <a href="https://arxiv.org/abs/1505.04597" rel="nofollow noopener">U-NET paper</a>, I am trying to implement custom weight maps to counter class imbalances.<br>
Below is the code for custom weight map-</p>
<pre><code class="lang-auto">from skimage.segmentation import find_boundaries

w0 = 10
sigma = 5

def make_weight_map(masks):
    """
    Generate the weight maps as specified in the UNet paper
    for a set of binary masks.
    
    Parameters
    ----------
    masks: array-like
        A 3D array of shape (n_masks, image_height, image_width),
        where each slice of the matrix along the 0th axis represents one binary mask.

    Returns
    -------
    array-like
        A 2D array of shape (image_height, image_width)
    
    """
    nrows, ncols = masks.shape[1:]
    masks = (masks &gt; 0).astype(int)
    distMap = np.zeros((nrows * ncols, masks.shape[0]))
    X1, Y1 = np.meshgrid(np.arange(nrows), np.arange(ncols))
    X1, Y1 = np.c_[X1.ravel(), Y1.ravel()].T
    for i, mask in enumerate(masks):
        # find the boundary of each mask,
        # compute the distance of each pixel from this boundary
        bounds = find_boundaries(mask, mode='inner')
        X2, Y2 = np.nonzero(bounds)
        xSum = (X2.reshape(-1, 1) - X1.reshape(1, -1)) ** 2
        ySum = (Y2.reshape(-1, 1) - Y1.reshape(1, -1)) ** 2
        distMap[:, i] = np.sqrt(xSum + ySum).min(axis=0)
    ix = np.arange(distMap.shape[0])
    if distMap.shape[1] == 1:
        d1 = distMap.ravel()
        border_loss_map = w0 * np.exp((-1 * (d1) ** 2) / (2 * (sigma ** 2)))
    else:
        if distMap.shape[1] == 2:
            d1_ix, d2_ix = np.argpartition(distMap, 1, axis=1)[:, :2].T
        else:
            d1_ix, d2_ix = np.argpartition(distMap, 2, axis=1)[:, :2].T
        d1 = distMap[ix, d1_ix]
        d2 = distMap[ix, d2_ix]
        border_loss_map = w0 * np.exp((-1 * (d1 + d2) ** 2) / (2 * (sigma ** 2)))
    xBLoss = np.zeros((nrows, ncols))
    xBLoss[X1, Y1] = border_loss_map
    # class weight map
    loss = np.zeros((nrows, ncols))
    w_1 = 1 - masks.sum() / loss.size
    w_0 = 1 - w_1
    loss[masks.sum(0) == 1] = w_1
    loss[masks.sum(0) == 0] = w_0
    ZZ = xBLoss + loss
    return ZZ
</code></pre>
<p>I am really not sure how to incorporate custom weights in <code>cross entropy loss</code> . I come across this example for <a href="https://jaidevd.github.io/posts/weighted-loss-functions-for-instance-segmentation/" rel="nofollow noopener">keras</a> but couldn’t find any source for pytorch</p>

username: KFrank
cooked: <p>Hello Rishav!</p>
<aside class="quote no-group quote-modified" data-username="Rishav_Sapahia" data-post="1" data-topic="56443" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/r/51bf81/48.png" class="avatar"> Rishav_Sapahia:</div>
<blockquote>
<p>…<br>
I am trying to implement custom weight maps to counter class imbalances.<br>
…<br>
I am really not sure how to incorporate custom weights in <code>cross entropy loss</code> .</p>
</blockquote>
</aside>
<p>Does the <code>weight</code> argument described in the <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.CrossEntropyLoss" rel="noopener nofollow ugc">documentation</a> for<br>
<code>CrossEntropyLoss</code> do what you need?</p>
<p>Best.</p>
<p>K. Frank</p>

username: Rishav_Sapahia
cooked: <p>Hi K. Frank,</p>
<p>The problem I am facing is to pass custom weight as an argument to <code>CrossEntropyLoss</code>, I have to pass masks as an argument like the the following method-<br>
<code>criterion = torch.nn.CrossEntropyLoss(make_weight_map(mask))</code><br>
which I guess is not the right way. I haven’t approached such kind of problem earlier so pretty clueless right now</p>

username: KFrank
cooked: <p>Hi Rishav!</p>
<aside class="quote no-group" data-username="Rishav_Sapahia" data-post="3" data-topic="56443" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/r/51bf81/48.png" class="avatar"> Rishav_Sapahia:</div>
<blockquote>
<p>The problem I am facing is to pass custom weight as an argument to <code>CrossEntropyLoss</code>, I have to pass masks as an argument like the the following method-<br>
<code>criterion = torch.nn.CrossEntropyLoss(make_weight_map(mask))</code><br>
which I guess is not the right way. I haven’t approached such kind of problem earlier so pretty clueless right now</p>
</blockquote>
</aside>
<p>I’m not clear on what you are asking.</p>
<p>Do you want the contribution of a given prediction-target<br>
pair to the loss to have a weight that depends (solely) on<br>
the class of the target?</p>
<p>If so, I believe that the <code>weight</code> argument of <code>CrossEntropyLoss</code><br>
does what you want.</p>
<p>If this isn’t the case, could you write out a simple equation that<br>
shows what you want the loss to be and how your “weights”<br>
enter into it?</p>
<p>You make it sound like the <code>weight</code> argument doesn’t do what<br>
you want, but you don’t say why.  If you could explain how<br>
<code>weight</code> doesn’t work, that might help make clear what it is<br>
you want to achieve.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Rishav_Sapahia
cooked: <p>Hi K. Frank,<br>
I want to apply the following operation-<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/f/8fd59e49947fae82f2d65e10a478384bdc9d6dd0.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/8fd59e49947fae82f2d65e10a478384bdc9d6dd0" title="25%20PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/f/8fd59e49947fae82f2d65e10a478384bdc9d6dd0_2_690x264.jpeg" alt="25%20PM" data-base62-sha1="kwpX6QhiAJFIMVC4BkoktuyJ9D2" width="690" height="264" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/f/8fd59e49947fae82f2d65e10a478384bdc9d6dd0_2_690x264.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/8/f/8fd59e49947fae82f2d65e10a478384bdc9d6dd0_2_1035x396.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/f/8fd59e49947fae82f2d65e10a478384bdc9d6dd0_2_1380x528.jpeg 2x" data-dominant-color="F0F0F0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">25%20PM</span><span class="informations">1726×662 338 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>According to my understanding, If I need to set custom weight using the below way <code>criterion = torch.nn.CrossEntropyLoss(make_weight_map(mask))</code>, I’ll be using the masks outside of data-loaders which I am really not sure will work or not?</p>

username: Rakshit_Kothari
cooked: <p>You could do the following:</p>
<pre><code class="lang-auto">criterion = torch.nn.CrossEntropy(reduction='none')
</code></pre>
<p>This ensures that the function will return a loss value for each element. You could then multiply the weights to each loss element.</p>
<pre><code class="lang-auto">gt # Ground truth, format torch.long
pd # Network output
W # per-element weighting based on the distance map from UNet
loss = criterion(pd, gt)
loss = W*loss # Ensure that weights are scaled appropriately
loss = torch.sum(loss.flatten(start_dim=1), axis=0) # Sums the loss per image
loss = torch.mean(loss) # Average across a batch
</code></pre>
<p>I’ve observed that a weight of 0 could provide complications. It might be a good idea to add a small <code>eps</code> factor to <code>W</code>.</p>

username: Rishav_Sapahia
cooked: <p>But I am converting masks to classes while feeding them into dataloader.</p>

username: Rakshit_Kothari
cooked: <p>Do you mean there is different mask for each class? If that’s the case then you’ll need to write your own cross-entropy loss function.</p>

username: Rishav_Sapahia
cooked: <p>No, Each mask has at max 4 classes(Red,blue,green and black), I did what you told apart from the last statement <code>loss = torch.mean(loss)</code> because I am calculating the average across batch using <code>train_loss = train_loss + ((1 / (batch_idx + 1)) * (loss.data - train_loss))</code>. I was encountering a error because my <code>make_weight_map(masks)</code> was expecting a numpy array instead of tensor and also I forgot to move tensor to cpu ,which I did now. but now, I guess my system has kind of stalled I guess.<br>
Edit:I got the memory error now though my ram size is 32 gb.</p>

username: Rishav_Sapahia
cooked: <p>okay, I raised my computation power to 61 GB Ram and 16 gb Tesla V100 gpu, and still got the memory error.I think the approach is wrong</p>

username: Rakshit_Kothari
cooked: <p>Hi Rishav its been awhile so I guess you got this working. Let me know if you need help fixing this issue.</p>