Imbalance in 3D binary segmentation

ORoberts: <p>Hello ! (I apologize if I do something wrong here, I am not used to asking question on forums)</p>
<p>I am working with a model to segment brain tumors from 3D brain MRIs, here doing binary segmentation.<br>
I created a U-net model taking the 3d object as whole. I do a bit of processing beforehand to read data approximately of size <code>1*150*220*220</code>, <code>1</code> because I only have one channel.<br>
The preprocessing trims black borders, giving images with variable size so I use a Fully Convolutional Network with batches of one images.</p>
<p>The problem is that the tumors represent a very small portion of the voxels: around 1/300, creating a great imbalance.<br>
I tried to use the <code>nn.BCEWithLogitLoss</code> as my loss function, allowing to weight the different classes, but it does not seem to be working correctly, the model giving me only negative answers (no voxel is a tumor)</p>
<pre><code class="lang-auto">weights = torch.tensor([300.])
loss_fn = torch.nn.BCEWithLogitsLoss(pos_weight=weights)
</code></pre>
<p>I tired to change the <code>300.</code> to different values, hopping to see a result, or at least a variation (<code>0.0001</code> or <code>50000</code>) but nothing ever appeared.</p>
<p>Am I using the function wrong ? Is there an other loss function that could do what I want it to do ?</p>
<p>Once again I apologize if this seems confused or easy, I have little experience using pytorch, so I do everything based on what I can find online.<br>
Thank you very much for your time and consideration.</p>

KFrank: <p>Hi Owen!</p>
<aside class="quote no-group quote-modified" data-username="ORoberts" data-post="1" data-topic="163923" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/oroberts/48/52185_2.png" class="avatar"> ORoberts:</div>
<blockquote>
<p>The problem is that the tumors represent a very small portion of the voxels: around 1/300, creating a great imbalance.</p>
<pre><code class="lang-auto">weights = torch.tensor([300.])
loss_fn = torch.nn.BCEWithLogitsLoss(pos_weight=weights)
</code></pre>
</blockquote>
</aside>
<p>This is indeed the correct approach to reweight your underrepresented<br>
positive class (assuming that your tumor voxels are labelled as your<br>
“positive” voxels).</p>
<p>The next step would be step-by-step debugging:</p>
<p>Read (or at least skim) through your code looking for typos or other bugs.</p>
<p>Make sure that you are <em>not</em> using a <code>sigmoid()</code> as a final “activation” in<br>
your network.  (You want to use the output of your final <code>out_channels = 1</code><br>
<code>Conv3d</code> layer as the predictions that are fed directly into your<br>
<code>BCEWithLogitsLoss</code> loss criterion.)</p>
<p>Check how you are converting the output of your network into “hard,”<br>
yes-no predictions for your voxels.</p>
<p>Look at your input images, output (prediction) images, and your label (mask)<br>
images to make sure everything looks okay.</p>
<p>Run a single forward pass and check the output.  Check the value of the<br>
loss function with and without <code>pos_weight</code>.</p>
<p>Run a single backward pass.  Do you get (reasonable) gradients?</p>
<p>Run a single optimizer step.  Do your weights change?</p>
<p>Run another forward pass on the same input image.  Does the output of<br>
your model change?</p>
<p>If all of this seems to be working, see if you can overfit a <em>small</em> number<br>
of input images by running many training epochs over and over again on<br>
the same input images.  Does the loss go down (even if slowly)?  Does<br>
the “probability” of a predicted positive voxel go up, even if slowly, and<br>
even if it doesn’t cross the threshold you’ve chosen to covert the “soft,”<br>
probabilistic prediction into a “hard,” yes-no prediction?</p>
<p>And so on …</p>
<p>Good luck!</p>
<p>K. Frank</p>

ORoberts: <p>Hello and thank you so much for helping !</p>
<p>I have done step-by-step debugging, analyzing the output of my model before any <code>sigmoid</code> function.</p>
<p>I came to see on a small set (5) of images, that the prediction of a default(not trained) have an average around 0.0.  Training the model a few epochs (10 at most) on thoses images results in the range going down, having no value above 0 (so no voxel counted as positive). The range of results has been for example [-100, -1].  But the iterations do change the output of my network so there is indeed a training</p>
<p>Changing the <code>pos_weight</code> using either <code>300</code> or <code>1</code> changes the loss of the examples by aroung 10^-4 on a default model (from <code>0.6872...</code> without the <code>pos_weight</code> to <code>0.6874...</code> with it)</p>
<p>I apologize, I am not sure how to interpret thoses results, if they should be considered ad “working” or not. There is indeed a training, but it goes in the wrong direction <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=12" title=":confused:" class="emoji" alt=":confused:" loading="lazy" width="20" height="20"></p>

KFrank: <p>Hi Owen!</p>
<aside class="quote no-group" data-username="ORoberts" data-post="3" data-topic="163923" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/oroberts/48/52185_2.png" class="avatar"> ORoberts:</div>
<blockquote>
<p>Changing the <code>pos_weight</code> using either <code>300</code> or <code>1</code> changes the loss of the examples by aroung 10^-4 on a default model (from <code>0.6872...</code> without the <code>pos_weight</code> to <code>0.6874...</code> with it)</p>
</blockquote>
</aside>
<p>Such values for <code>BCEWithLogitsLoss</code> in the context you describe look a<br>
little odd to me.</p>
<p>Could you tell me in words what the last three operations your network<br>
performs on the values it outputs as its predictions are?  Please also then<br>
post the last three layers of your network.</p>
<p>Please post the line where you instantiate your optimizer and the<br>
section of your code from <code>predictions = model (input)</code> through<br>
<code>optimizer.step()</code>.</p>
<p>Print out the shape, type, min, and max of one of your 3d input images, the<br>
shape, type, min, and max of the output your network predicts for it, together<br>
with the shape, type, min, and max of its corresponding ground-truth mask.</p>
<p>Best.</p>
<p>K. Frank</p>

ORoberts: <p>Hello K. Frank !</p>
<p>The last three operations are a LeakyReLU and 2 3D Convolutions with a kernel of <code>size=5</code><br>
such ass</p>
<pre><code class="lang-auto">nn.LeakyReLU(),
nn.Conv3d(in_channels=16, out_channels=16, kernel_size=5, padding=2),
nn.Conv3d(in_channels=16, out_channels=1, kernel_size=5, padding=2),
</code></pre>
<p>as for the code itself</p>
<pre><code class="lang-auto">lr=1e-3
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
weights = torch.tensor([300.]).to(device)
loss_fn = torch.nn.BCEWithLogitsLoss(pos_weight=weights, reduction='mean')
</code></pre>
<pre><code class="lang-auto">def epoch_train_all(model: Baseline, optim, loss_fn, device, batch_size:int = 5, path: str = "./data/train/"):
    optim.zero_grad()
    avg_loss = 0.0
    nb_files = 0
    model.to(device)
    for id, dm in get_files_it(path):
        nb_files += 1
        print(f"\r\tTRAIN file nb{nb_files:03d}, id:{id}", end=" ")
        # transferring data from cache to device
        data = dm.data.to(device)
       


        pred = model(data)
        del data
        mask = dm.mask.to(device)
        loss = loss_fn(pred, mask)
        del pred
        del mask
        # once the data has been used, free it to free some space on the device
        loss.backward()
        avg_loss += float(loss)
        del loss
        if nb_files % batch_size == 0:
          
            print("optim !", end='')
            optim.step()
            optim.zero_grad()
            return




    optim.step()
    optim.zero_grad()
    avg_loss /= nb_files
    print(f"avg. train loss: {avg_loss}")
    return avg_loss

</code></pre>
<p>I have isolated in the middle the part that you asked for.<br>
I will add a few precisions:</p>
<ul>
<li>I use gradient accumulation, so my batch size can control the frequency of my backward pass</li>
<li>there is here a <code>return</code> after the 5 items in my batch to allow fitting on only the first 5 elements of my data (I have disabled the shuffling)</li>
<li>
<code>get_files_it(path):</code> is an iterator I have made which returns my processed data in a class I have created</li>
<li>I delete the elements once they are not used anymore to save some memory</li>
</ul>
<p>As for the statistics:<br>
data: input data<br>
prediction: the output value of my network<br>
pred w/ sig: the previous prediction passed into a sigmoid (not in the network because working bit BCEWithLogitsLoss)<br>
ground truth: the expected mask</p>
<p>1st epoch:</p>
<pre><code class="lang-auto">data:        	shape=torch.Size([1, 1, 248, 88, 176]), type: torch.cuda.FloatTensor, range:[0.0; 1.0]
prediction: 	shape=torch.Size([1, 1, 248, 88, 176]), type: torch.cuda.FloatTensor, range:[-0.016314834356307983; 0.0024049831554293633]
pred w/ sig:	shape=torch.Size([1, 1, 248, 88, 176]), type: torch.cuda.FloatTensor, range:[0.49592143297195435; 0.5006012320518494]
ground truth:	shape=torch.Size([1, 1, 248, 88, 176]), type: torch.cuda.FloatTensor, range:[0.0; 0.0021929824724793434]
</code></pre>
<p>10th epoch</p>
<pre><code class="lang-auto">data:       	shape=torch.Size([1, 1, 248, 88, 176]), type: torch.cuda.FloatTensor, range:[0.0; 1.0]
prediction: 	shape=torch.Size([1, 1, 248, 88, 176]), type: torch.cuda.FloatTensor, range:[-269.3184814453125; -1.7151740789413452]
pred w/ sig:	shape=torch.Size([1, 1, 248, 88, 176]), type: torch.cuda.FloatTensor, range:[0.0; 0.15249381959438324]
ground truth:	shape=torch.Size([1, 1, 248, 88, 176]), type: torch.cuda.FloatTensor, range:[0.0; 0.0021929824724793434]
</code></pre>
<p>Thank you very much for your time and help, I appreciate <img src="https://discuss.pytorch.org/images/emoji/apple/relaxed.png?v=12" title=":relaxed:" class="emoji" alt=":relaxed:" loading="lazy" width="20" height="20"></p>

KFrank: <p>Hi Owen!</p>
<p>Thank you for the details.</p>
<aside class="quote no-group quote-modified" data-username="ORoberts" data-post="5" data-topic="163923" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/oroberts/48/52185_2.png" class="avatar"> ORoberts:</div>
<blockquote>
<pre><code class="lang-auto">        loss = loss_fn(pred, mask)
</code></pre>
<p>…<br>
ground truth: the expected mask<br>
…</p>
<pre><code class="lang-auto">ground truth:	shape=torch.Size([1, 1, 248, 88, 176]), type: torch.cuda.FloatTensor, range:[0.0; 0.0021929824724793434]
</code></pre>
</blockquote>
</aside>
<p>Assuming that “ground truth” is the <code>mask</code> that you pass into <code>loss_fn()</code>,<br>
this is likely your problem.  <code>BCEWithLogitsLoss</code> expects it <code>target</code> (your<br>
<code>mask</code>) to be the probability of the corresponding item to be in the “positive”<br>
class.  Based on the range you posted, some of your <code>mask</code> values are<br>
<code>0.0</code> which means definitely in the “negative” class (that is, 0% probability<br>
of being in the “positive” class).  However, your largest <code>mask</code> value is<br>
<code>0.00219</code> which means that the corresponding item has only a <code>0.219%</code><br>
probability of being in the “positive” class which means that it is <em>almost<br>
certainly</em> in the “negative” class (in the “negative” class with a probability<br>
of <code>99.78%</code>).</p>
<p>Therefore your ground-truth <code>mask</code> is saying that <em>all</em> of your voxels are<br>
either certainly “negative” or almost certainly “negative.”  So you are<br>
(successfully) training your network to always predict “negative” voxels<br>
(regardless of the value of <code>pos_weight</code>, because, according to <code>mask</code>,<br>
you don’t have any positive voxels).</p>
<p>I don’t know where your <code>mask</code> comes from, but could it be that <code>0.0</code> means<br>
“negative” while <em>any</em> non-zero value means “positive?”  You would want<br>
“positive” voxels to be labelled with a <code>value</code> equal to, or at least reasonably<br>
close to <code>1.0</code>.</p>
<p>Regardless of the scale of the values, is <code>mask</code> supposed to contain hard<br>
“negative”-or-“positive” labels, or are the values to be understood as soft,<br>
probabilistic labels?</p>
<p>Best.</p>
<p>K. Frank</p>

ORoberts: <p>Hello KFrank,</p>
<p>I feel so stupid for not seeing this.<br>
I was applying my normalization wrong on my masks (they come with 0 as the negative value and 2 as the positive value). I was using the normalization I apply on my data, giving stupidly low results. I am now trying with the corrected data.</p>
<p>The masks are boleans, they should be either yes-or-no values per voxels.</p>
<p>Thank you very much,<br>
Owen</p>