is invariant 3D Unet even possible?

username: kfir7755
cooked: <p>i have a project in torch to create a 3D semantic segmentation with 3D MRI data stored in NRRD files (that can be converted to 3D tensors). The images are in different but pretty similar shapes.</p>
<p>I tried to create something similar to 3D Unet but in invariant version using <strong>torch.nn.Conv3d</strong> in the encoder blocks and <strong>torch.nn.ConvTranspose3d</strong> in decoder blocks but the model wont return the same shape unless all 3 dimensions of the image shape are divided without remainders in 16 (there are 4 encoder blocks with stride of 2).<br>
Is there a way to fix it?</p>
<p>Also possible solution I thought of that is possible if invariant 3D Unet is not possible was padding the image with 0 to reach a shape that is divisible by 16. I am sure that it’s not optimal so if there are other ideas I would love to know them <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>
<p>Thanks to everyone in advance</p>

username: KFrank
cooked: <p>Hi kfir!</p>
<aside class="quote no-group" data-username="kfir7755" data-post="1" data-topic="192179" 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/k/e79b87/48.png" class="avatar"> kfir7755:</div>
<blockquote>
<p>I tried to create something similar to 3D Unet</p>
</blockquote>
</aside>
<p>Preliminary comment:</p>
<p>The lengths of your three dimensions (height, width, depth) of your 3D<br>
U-Net don’t mix as your 3D image passes through the U-Net in the<br>
following sense:  The “valid” length of, say, the height dimension (which<br>
you say in your case is “divided without remainders in 16”) doesn’t<br>
depend on the values of the other two dimensions.  Each of the three<br>
dimensions can be valid or invalid independently.</p>
<p>So I don’t think the fact that you are working with a 3D U-Net is relevant<br>
to the question I think your asking – your question applies equally to 2D<br>
and 3D U-Nets.</p>
<aside class="quote no-group">
<blockquote>
<p>in invariant version</p>
</blockquote>
</aside>
<p>I’m not sure what you mean by an “invariant U-Net.”  Are you asking that<br>
the output image have the same shape as the input image (not counting<br>
the channels dimension)?</p>
<p>If so, I don’t think you want this.  As your image passes through a U-Net,<br>
its edges get nibbled away by the convolutions.  Padding the convolutions<br>
so that the image doesn’t shrink pollutes somewhat the U-Net logic.  Think<br>
about how padding the convolutions would be injecting “noise” at the<br>
image edges at each convolution stage.</p>
<p>I highly recommend carefully reading the <a href="https://arxiv.org/abs/1505.04597" rel="noopener nofollow ugc">original U-Net paper,</a> paying<br>
attention to how the shape of the image shrinks as it flows through the<br>
U-Net.  Note that the original paper does not pad its convolutions.</p>
<aside class="quote no-group">
<blockquote>
<p>the model wont return the same shape unless all 3 dimensions of the image shape are divided without remainders in 16</p>
</blockquote>
</aside>
<p>The condition for “valid” input-image dimensions is more nuanced for<br>
the original U-Net (and I believe this will be true for any U-Net with<br>
optimally-clean convolution / no-padding logic).  Look at the U-Net<br>
diagram in the original paper.  It shows how the image shrinks and<br>
from it you can deduce the rules for valid input dimensions.</p>
<aside class="quote no-group">
<blockquote>
<p>Is there a way to fix it?</p>
</blockquote>
</aside>
<p>If by “fix it,” you mean have your U-Net output an image whose size<br>
matches that of its input, you can achieve this by appropriately padding<br>
each of the convolutions along the way.  But, as discussed above, you<br>
don’t really want to do this.</p>
<aside class="quote no-group">
<blockquote>
<p>I thought of that is possible if invariant 3D Unet is not possible was padding the image with 0 to reach a shape that is divisible by 16.</p>
</blockquote>
</aside>
<p>My recommendation is to reflection-pad the input image up to a size<br>
that has valid dimensions (which I think will be more complicated than<br>
“divisible by 16”).  Your image edges will have a somewhat different<br>
character than the interior, but this is already the reality before you pad,<br>
and I think that reflection padding is the “gentlest” way the let your model<br>
learn to segment pixels near the edges.</p>
<p>Best.</p>
<p>K. Frank</p>

username: kfir7755
cooked: <p>Hi Frank, thanks for your comment.</p>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="192179">
<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/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>your question applies equally to 2D<br>
and 3D U-Nets.</p>
</blockquote>
</aside>
<p>You are definitely right about this one, I didn’t want to get options like resizing the image to a shape that is divisible by 16, because it’s only possible with 2D images and not 3D so I focused about 3D but creating invariant 2D Unet is the same problem as creating invariant 3D Unet for me.</p>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="192179">
<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/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>I’m not sure what you mean by an “invariant U-Net.” Are you asking that<br>
the output image have the same shape as the input image (not counting<br>
the channels dimension)?</p>
</blockquote>
</aside>
<p>Yes, this is exactly what I meant</p>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="192179">
<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/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>I highly recommend carefully reading the <a href="https://arxiv.org/abs/1505.04597" rel="noopener nofollow ugc">original U-Net paper,</a> paying<br>
attention to how the shape of the image shrinks as it flows through the<br>
U-Net. Note that the original paper does not pad its convolutions.</p>
</blockquote>
</aside>
<p>I never noticed that in the original U-net padding wasn’t used <img src="https://discuss.pytorch.org/images/emoji/apple/hushed.png?v=12" title=":hushed:" class="emoji" alt=":hushed:" loading="lazy" width="20" height="20"><br>
I’ll read it carefully as you suggested I believe it is important, although in many researches using similar architecture, padding was used some of the time.</p>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="192179">
<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/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>My recommendation is to reflection-pad the input image up to a size<br>
that has valid dimensions (which I think will be more complicated than<br>
“divisible by 16”). Your image edges will have a somewhat different<br>
character than the interior, but this is already the reality before you pad,<br>
and I think that reflection padding is the “gentlest” way the let your model<br>
learn to segment pixels near the edges.</p>
</blockquote>
</aside>
<p>Didn’t think about this option, will definitely try since It sounds very reasonable.<br>
Again thank you for your highly detailed answer Frank, really helped me!</p>

username: KFrank
cooked: <p>Hi kfir!</p>
<aside class="quote no-group" data-username="kfir7755" data-post="3" data-topic="192179" 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/k/e79b87/48.png" class="avatar"> kfir7755:</div>
<blockquote>
<p>resizing the image to a shape that is divisible by 16, because it’s only possible with 2D images and not 3D</p>
</blockquote>
</aside>
<p>Just to be clear, any resizing would happen on a dimension-by-dimension<br>
basis, so there is no substantive difference between 2D and 3D.</p>
<aside class="quote no-group">
<blockquote>
<p>although in many researches using similar architecture, padding was used some of the time.</p>
</blockquote>
</aside>
<p>I think that many people are not as insightful and / or careful as the authors<br>
of the original U-Net paper.  They really did a quite nuanced job of it and I<br>
think that not padding internally is the right way to go.</p>
<p>It can be a real pain in the neck to prepare batches of images of varying<br>
size to have valid dimensions for input into a U-Net, so there is a certain<br>
motivation to cut corners, rather than do it right.</p>
<p>Best.</p>
<p>K. Frank</p>

username: J_Johnson
cooked: <aside class="quote no-group" data-username="kfir7755" data-post="1" data-topic="192179">
<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/k/e79b87/48.png" class="avatar"> kfir7755:</div>
<blockquote>
<p>padding the image with 0 to reach a shape that is divisible by 16</p>
</blockquote>
</aside>
<p>If you’re just padding the input image and corresponding targets during preprocessing to be a standardized size, that should be fine. For example, using:<br>
<a href="https://pytorch.org/vision/stable/generated/torchvision.transforms.v2.Pad.html#torchvision.transforms.v2.Pad" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/vision/stable/generated/torchvision.transforms.v2.Pad.html#torchvision.transforms.v2.Pad</a></p>
<p>I think what <a class="mention" href="/u/kfrank">@KFrank</a> is getting at is you do not want to add padding into any of the convolution operations in the UNet.</p>
<aside class="quote no-group" data-username="kfir7755" data-post="1" data-topic="192179">
<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/k/e79b87/48.png" class="avatar"> kfir7755:</div>
<blockquote>
<p>3D Unet but in invariant version</p>
</blockquote>
</aside>
<p>Regarding an invariant UNet, it’s not really possible/practical, especially given the problem you intend to use it for. As far as I can tell, you have 3D MRI scans, and corresponding targets of identical size. And you want to train the model to precisely identify what pixels contain a particular target.</p>
<p>If you used <a href="https://pytorch.org/docs/stable/generated/torch.nn.AdaptiveAvgPool3d.html" rel="noopener nofollow ugc">nn.AdaptiveAvgPool3d</a> within the encoder of the model, you’re going to have a major problem with getting the decoder to give the correct output size that matches your targets, because the size information will be lost.</p>
<p>So any resizing or image padding should be handled in your image preprocessing and applied the same to your targets so they match pixel for pixel.</p>

username: kfir7755
cooked: <p>Hi, I read the original U-net paper and now what you <a class="mention" href="/u/j_johnson">@J_Johnson</a> <a class="mention" href="/u/kfrank">@KFrank</a> are saying is much more clear and makes sense, I think I will remove padding from the network itself as suggested by you and will use mirror-padding in the pre processing for the images.<br>
Thanks a lot really appreciate the help and explanations.</p>
<p>Kfir.</p>

username: kfir7755
cooked: <p>for anyone who will encounter this post later on, apparently there is a function in the libraray monai - monai.transforms.DivisiblePad which does the padding for this case</p>