U-Net collapse during Training

username: twangmagister
cooked: <p>Hey all,</p>
<p>I am currently training a U-Net to predict a noise in a given image (diffusion) and encounter some weird issues. Mainly, the loss spikes at some point in time and from then stays roughly the same, shown in this image:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/f/af01d34164e311ea673752a55059a468c3f80a27.png" data-download-href="https://discuss.pytorch.org/uploads/default/af01d34164e311ea673752a55059a468c3f80a27" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/f/af01d34164e311ea673752a55059a468c3f80a27_2_517x141.png" alt="image" data-base62-sha1="oYbt0mFUBkl64i4kpQZj65rY0F9" width="517" height="141" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/f/af01d34164e311ea673752a55059a468c3f80a27_2_517x141.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/f/af01d34164e311ea673752a55059a468c3f80a27_2_775x211.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/f/af01d34164e311ea673752a55059a468c3f80a27_2_1034x282.png 2x" data-dominant-color="353332"><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">image</span><span class="informations">2053×565 45.6 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>I use AdamW and a simple MSE between the noise and the output of the U-Net. For my experiment now I use a batch size of 1, but the issue also appears with larger batch sizes. Regarding this I have two specific questions:</p>
<p><strong>1. What is causing the spike? (How can it be avoided?)</strong><br>
<strong>2. Why does the model stop learning after the spike?</strong></p>
<p>I tried to fix the issue with gradient clipping and it worked, but the results turn out to be bad. So I want to fix the problem at its “root”. I ran some experiments, some insights are shown below.</p>
<p>The moving averages of Adam are also affected:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/f/5f11994ee64bba2a7fb760cd5b627ad09245cf36.png" data-download-href="https://discuss.pytorch.org/uploads/default/5f11994ee64bba2a7fb760cd5b627ad09245cf36" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/f/5f11994ee64bba2a7fb760cd5b627ad09245cf36_2_517x285.png" alt="image" data-base62-sha1="dz1429ACTZxUy81x2SA40oCM0KO" width="517" height="285" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/f/5f11994ee64bba2a7fb760cd5b627ad09245cf36_2_517x285.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/f/5f11994ee64bba2a7fb760cd5b627ad09245cf36_2_775x427.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/f/5f11994ee64bba2a7fb760cd5b627ad09245cf36_2_1034x570.png 2x" data-dominant-color="3B3432"><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">image</span><span class="informations">2054×1133 112 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>I ran multiple (9) experiments with the same random seed (pytorch_lightnings seed_everyhting) and I noticed some weird behaviour; up to a certain step (14) the behaviour is always almost identical, from then, losses, gradients etc. start to differ:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/e/aec68450eb1ca22dbe0d7e59d8ac0ef9296fac0e.png" data-download-href="https://discuss.pytorch.org/uploads/default/aec68450eb1ca22dbe0d7e59d8ac0ef9296fac0e" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/e/aec68450eb1ca22dbe0d7e59d8ac0ef9296fac0e_2_517x227.png" alt="image" data-base62-sha1="oW8oODcayXqlHs6VMNrZqK0GqUC" width="517" height="227" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/e/aec68450eb1ca22dbe0d7e59d8ac0ef9296fac0e_2_517x227.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/e/aec68450eb1ca22dbe0d7e59d8ac0ef9296fac0e_2_775x340.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/e/aec68450eb1ca22dbe0d7e59d8ac0ef9296fac0e_2_1034x454.png 2x" data-dominant-color="343334"><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">image</span><span class="informations">1337×588 78.4 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><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/4/447e4b4af4715d950094e8cc7c5cbe7c9046f7fc.png" data-download-href="https://discuss.pytorch.org/uploads/default/447e4b4af4715d950094e8cc7c5cbe7c9046f7fc" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/447e4b4af4715d950094e8cc7c5cbe7c9046f7fc_2_517x186.png" alt="image" data-base62-sha1="9LV3dq7nQXH0SFECQJxRXJlk4Uc" width="517" height="186" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/447e4b4af4715d950094e8cc7c5cbe7c9046f7fc_2_517x186.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/447e4b4af4715d950094e8cc7c5cbe7c9046f7fc_2_775x279.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/4/4/447e4b4af4715d950094e8cc7c5cbe7c9046f7fc_2_1034x372.png 2x" data-dominant-color="363536"><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">image</span><span class="informations">1640×590 146 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><strong>Any idea what happens at step 14 or 15 that changes the model outcome?</strong></p>
<p>Left side Model Output, Right Side GT (before collapse):<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/b/5/b5274c94d947aaa461915d61f92b491ab6276d84.png" alt="image" data-base62-sha1="pQyCrZaicN3brSJh3ytWVuuWmaM" width="392" height="205"><br>
Left side Model Output, Right Side GT (“during” collapse):<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/e/f/ef4b511d84d339004325cab327f0e0eea163d4ad.jpeg" alt="image" data-base62-sha1="y8TtPbIw4z90K0azGtySAVV7qJD" width="412" height="210"><br>
Left side Model Output, Right Side GT (after collapse):<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/b/9/b941e9e4e72818d3b448c3eb61b57e238f473772.png" alt="image" data-base62-sha1="qqRybu3JCH8KQRKF6idQz5tXznY" width="417" height="213"></p>
<p><strong>I also noticed that this only happens for large models and not smaller ones.</strong></p>
<p>Specs:<br>
Ubuntu 20.04<br>
RTX 3090<br>
CUDA 11.7<br>
Python 3.7.13<br>
Pytorch 1.11.0<br>
Pytorch Lightning 1.6.1 (used for training)</p>
<p>I am thankfull for any suggestions regarding the cause of the issues or suggestions on how I can continue my analysis. Thanks!</p>

username: Xern
cooked: <p>Hey, i have the same. Did you find answer? Gradient clipping doesn’t help in my case.</p>