Vram maxed (even when changing the batch size?)

AladdinPerzon: <p>I’m playing around with ProGAN and getting some weird behavior related to vram usage, especially at larger image size it seems to always take very close to the maximum amount on my gpu, even when lowering the batch size. This leads to it being quite unstable as I’m getting OOM error messages sometimes.</p>
<p>I believe there might be something related to my train function and that I might be doing something incorrectly here. Perhaps someone could take a look if something looks off. Specifically I’m a bit unsure of how to use the scaler when it comes to several optimizers</p>
<pre><code class="lang-auto">def train_fn(
    critic,
    gen,
    loader,
    dataset,
    step,
    alpha,
    opt_critic,
    opt_gen,
    tensorboard_step,
    writer,
    scaler_gen,
    scaler_critic,
):
    loop = tqdm(loader, leave=True)

    for batch_idx, (real, _) in enumerate(loop):
        real = real.to(config.DEVICE)
        cur_batch_size = real.shape[0]

        for _ in range(config.CRITIC_ITERATIONS):
            noise = torch.randn(cur_batch_size, config.Z_DIM, 1, 1).to(config.DEVICE)

            with torch.cuda.amp.autocast():
                fake = gen(noise, alpha, step)
                critic_real = critic(real, alpha, step).reshape(-1)
                critic_fake = critic(fake, alpha, step).reshape(-1)
                gp = gradient_penalty(critic, real, fake, alpha, step, device=config.DEVICE)
                loss_critic = (
                    -(torch.mean(critic_real) - torch.mean(critic_fake))
                    + config.LAMBDA_GP * gp
                )

            opt_critic.zero_grad()
            scaler_critic.scale(loss_critic).backward(retain_graph=True)
            scaler_critic.step(opt_critic)
            scaler_critic.update()

        # Train Generator: max E[critic(gen_fake)] &lt;-&gt; min -E[critic(gen_fake)]
        with torch.cuda.amp.autocast():
            fake = gen(noise, alpha, step)
            gen_fake = critic(fake, alpha, step).reshape(-1)
            loss_gen = -torch.mean(gen_fake)

        opt_gen.zero_grad()
        scaler_gen.scale(loss_gen).backward()
        scaler_gen.step(opt_gen)
        scaler_gen.update()

        # Update alpha and ensure less than 1
        alpha += cur_batch_size / (
            (config.PROGRESSIVE_EPOCHS[step]*0.5) * len(dataset) # - step
        )
        alpha = min(alpha, 1)

</code></pre>
<p>This behavior is however consistent in FP32 and FP16 so I don’t know what could be wrong.</p>

ptrblck: <aside class="quote no-group" data-username="AladdinPerzon" data-post="1" data-topic="112789">
<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/aladdinperzon/48/13045_2.png" class="avatar"> AladdinPerzon:</div>
<blockquote>
<p>This leads to it being quite unstable as I’m getting OOM error messages sometimes.</p>
</blockquote>
</aside>
<p>Are you getting the OOMs in the same iterations with the large and small batch sizes?<br>
If so, are you changing the input shapes or are you using static shapes?<br>
If you lower the batch size e.g. to 1, are you still seeing OOMs?<br>
Did you check the GPU memory usage via <code>nvidia-smi</code> or with PyTorch methods (e.g. <code>torch.cuda.memory_summary()</code>)?<br>
Note that the former would return all allocated memory, i.e. the CUDA context, the used and allocated memory as well as the cache in PyTorch, and the memory usage of all other applications.</p>

AladdinPerzon: <p>So I don’t know how but I don’t seem to be getting this problem anymore. Normally though if I change the batch size I can see the allocated memory decrease but in this situation it was always maxed. I have encountered this behavior in TensorFlow but never before in Pytorch. Sometimes (but rarely) this would cause an OOM randomly during an epoch (was using static shapes).</p>
<p>Interestingly I also upgraded my Pytorch version which resulted in VRAM halved and an epoch speed up by ~33%, which is pretty insane…</p>
<p>Edit: The performance improvements seems to be specific to the 30-series GPU (was just slow to notice this):</p><aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/47039" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/47039" target="_blank" rel="noopener nofollow ugc">Convolution operations are extremely slow on RTX 30 series GPU</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-10-29" data-time="04:37:03" data-timezone="UTC">04:37AM - 29 Oct 20 UTC</span>
      </div>

        <div class="date">
          closed <span class="discourse-local-date" data-format="ll" data-date="2020-10-29" data-time="14:51:56" data-timezone="UTC">02:51PM - 29 Oct 20 UTC</span>
        </div>

      <div class="user">
        <a href="https://github.com/P3n9W31" target="_blank" rel="noopener nofollow ugc">
          <img alt="P3n9W31" src="https://avatars.githubusercontent.com/u/25482246?v=4" class="onebox-avatar-inline" width="20" height="20">
          P3n9W31
        </a>
      </div>
    </div>
  </div>
</div>

<div class="github-row">
  <p class="github-content">🐛 Bug
I'm using PyTorch on an RTX 3090 GPU, and trying to train different types of networks (e.g., CNN, RNN, Transformer),...</p>
</div>

<div class="labels">
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>