How to train the original U-Net model?

dalifreire: <p>Hi guys!</p>
<p>I’m trying to implement and train the original U-Net model, but I’m stuck in when I’m trying to train the model using the ISBI Challenge Dataset.</p>
<p>According with the original U-Net model, the network outputs an image with 2 channels and size of 388 x 388. So, my data loader for training generates a tensor with size of <em>[batch, channels=1, width=572, height=572]</em> for the input images and <em>[batch, channels=2, width=388, width=388]</em> for target/output images.</p>
<p>My problem actually is that when I’m trying to use the nn.CrossEntropyLoss() the following error is raised:</p>
<blockquote>
<p>RuntimeError: invalid argument 3: only batches of spatial targets supported (3D tensors) but got targets of dimension: 4 at /opt/conda/conda-bld/pytorch_1556653099582/work/aten/src/THNN/generic/SpatialClassNLLCriterion.c:59</p>
</blockquote>
<p>I’m just starting with PyTorch (newbie here)… so, I’ll really appreciate if someone could help me to overcome this problem.</p>
<p>The sourcecode is available on GitHub: <a href="https://github.com/dalifreire/cnn_unet_pytorch" rel="noopener nofollow ugc">https://github.com/dalifreire/cnn_unet_pytorch</a><br>
<a href="https://github.com/dalifreire/cnn_unet_pytorch/blob/master/unet_pytorch.ipynb" rel="noopener nofollow ugc">https://github.com/dalifreire/cnn_unet_pytorch/blob/master/unet_pytorch.ipynb</a></p>
<p>Best regards!</p>

simon_thom: <p>How many classes do you have?</p>
<p>In that you say it outputs two channels, I’d assume that you’re making a binary segmentation.</p>
<p>According to <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.CrossEntropyLoss" rel="nofollow noopener">https://pytorch.org/docs/stable/nn.html#torch.nn.CrossEntropyLoss</a>, your target should be just a 2D map where the pixels has the value of the appropriate class, i.e. a binary mask with 0’s and 1’s. Hope this makes sense.</p>

dalifreire: <p>Hi Simon,</p>
<p>First of all, thanks for your reply!<br>
I’m trying to implement the same thing presented in the <a href="https://arxiv.org/pdf/1505.04597.pdf" rel="nofollow noopener">original paper</a>… i.e. <strong>two classes</strong> used for cell segmentation. So the masks contains only pixels with values ‘0’ and ‘1’.</p>
<p><a href="https://github.com/dalifreire/cnn_unet_pytorch/blob/master/isbi_dataset.py#L103" rel="nofollow noopener">Here</a> are the point in the sourcecode were I’m creating the masks.</p>
<p>Kind regards</p>

simon_thom: <p>Hi again</p>
<p>I’ve recently done something similar <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>I think for the nn.CrossEntropy to work,</p>
<p>your mask/target should have the shape:<br>
torch.Size([batch_size, 388, 388])</p>
<p>while your output from the unet should have the shape:<br>
torch.Size([batch_size, 2, 388, 388])</p>
<p>I’d guess a simple fix is to change this line:</p>
<pre><code class="lang-auto"># Crop the mask to the desired output size
mask = transforms.CenterCrop(img_output_size)(mask)
</code></pre>
<p>so you don’t get the second channel which the output has <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>BR</p>

dalifreire: <p>Simon,</p>
<p>Thanks a lot!<br>
I just remove the channel dimension from my masks and everything works well… <a href="https://github.com/dalifreire/cnn_unet_pytorch/blob/master/isbi_dataset.py#L49" rel="nofollow noopener">now I’m generating masks with the shape</a> <em>[width=388, height=388]</em>.</p>
<p>After that, I’m working with input images (X), target masks (y) and predicted output masks (y_hat) as follow:</p>
<pre><code class="lang-auto">X     --&gt; torch.Size([10, 1, 572, 572])
y     --&gt; torch.Size([10, 388, 388])
y_hat --&gt; torch.Size([10, 2, 388, 388])
</code></pre>
<p>But, I don’t understand why target masks (y) and predicted masks (y_hat) must have different shapes? It’s so weird for me…</p>
<p><a href="https://github.com/dalifreire/cnn_unet_pytorch/blob/master/unet_pytorch.ipynb" rel="nofollow noopener">Here</a> are my jupyter notebook that I’m using in my tests.</p>