Calling loss.item() is very slow

NightRain: <p>When calling loss.item() in the “standard” training loop (e.g. <a href="https://github.com/pytorch/examples/blob/master/mnist/main.py" rel="noopener nofollow ugc">here</a>), the operation takes a lot of time (~2 seconds). Is there a way to reduce this bottleneck? I understand that the problem might be the loss being in the GPU…</p>
<p>Thanks!</p>

ptrblck: <p>The operation itself should not be slow, but is synchronizing the code, if you are using the GPU.<br>
I.e. since CUDA operations are executed asynchronously, the timing will be accumulated in the next synchronizing operation, which is <code>tensor.item()</code>, since it needs to wait for the GPU calculation to finish in order to grab the value and push it to the CPU.</p>

NightRain: <p>Thanks, is there any way to speed this up (by e.g. printing loss asynchronously when the operation is finished instead of waiting for the operation to finish)?</p>

fgauthier: <p>Did you try with a num_workers=0. By doing that I think you will eliminate the need of synchronizing the code so it might be faster. For me it was the problem, but I had to do a lot of guess before finding the solution. It might not be your solution but maybe it will work! (num_workers is a parameter of your dataloader) (I have only been working with PyTorch for 6 months so I might do an error!)</p>

ptrblck: <p>No, there isn’t a way to speed this up, as you need to wait for the GPU to finish the <code>loss</code> calculation before you can print this value. In case the GPU isn’t finished yet, the <code>print</code> statement will wait for it.</p>

NightRain: <p>If you do not print the value, but instead use a SummaryWriter for Tensorboard, is there a way to add the value to the writer in an asynchronous way (e.g. in a separate thread, when the GPU has finished fetching the loss, it will be added to the writer corresponding to the unique step specified in the add_scalar argument)?<br>
Have not seen this anywhere but would be faster especially when training times are high.</p>

steveazzolin: <p>If your goal is to print, for each epoch, the cumulative loss in order to analyze in TensorBoard/WANDB the performances of the model you might do something like this:</p>
<pre><code class="lang-auto">cumulative_loss = 0
with torch.set_grad_enabled(train):
    for batch_idx, (inputs, targets) in enumerate(data_loader):
      inputs, targets = inputs.to(device, non_blocking=True), targets.to(device, non_blocking=True) 
      outputs = net(inputs) 
      loss = cost_function(outputs, targets) 
      loss2 = cost_function(outputs, targets)
      if train:
        loss.backward() # Backward pass
        optimizer.step() # Update parameters
        optimizer.zero_grad() # Resets the gradients

      #cumulative_loss += loss.item()  very slow
      cumulative_loss += loss2   #should not need synch

    cumulative_loss = cumulative_loss.item() #synch once per epoch instead of once per batch
    logging_loss_epoch(cumulative_loss)
</code></pre>
<p>After some experiments to me it seems that the performances are much better.</p>

ptrblck: <p>I’m not sure why you are calculating the loss twice, but note that appending the loss tensor without detaching it will increase the memory usage, since the entire computation graph is attached, and could yield out of memory errors (which users often call a “memory leak”), so you might want to call <code>.detach()</code> on <code>loss2</code>.</p>

steveazzolin: <p>The reason is that I noticed that recomputing the loss is anyway faster than calling <code>.item()</code> on the loss. But actually you’re right. A cleaner version would replace the recomputation of the loss with simply <code>.detach()</code>, thanks.</p>

saypal: <p>Doesn’t calling loss.backward() detach the graph? I am getting the exact same time whether or not I call loss.detach(). (Also in my case, the time is not too different from just doing loss.item() every time.)</p>

ptrblck: <aside class="quote no-group" data-username="saypal" data-post="10" data-topic="99774">
<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/saypal/48/44821_2.png" class="avatar"> saypal:</div>
<blockquote>
<p>Doesn’t calling loss.backward() detach the graph?</p>
</blockquote>
</aside>
<p>No, <code>loss.backward()</code> will free the intermediate activations, but won’t detach the tensor from the computation graph. You might not see the large memory accumulation, but the proper approach is to detach the loss before storing it for printing purposes.</p>
<aside class="quote no-group" data-username="saypal" data-post="10" data-topic="99774">
<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/saypal/48/44821_2.png" class="avatar"> saypal:</div>
<blockquote>
<p>I am getting the exact same time whether or not I call loss.detach().</p>
</blockquote>
</aside>
<p>I don’t think detaching the loss would result in a large speed difference, but would avoid accumulating memory.</p>
<aside class="quote no-group" data-username="saypal" data-post="10" data-topic="99774">
<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/saypal/48/44821_2.png" class="avatar"> saypal:</div>
<blockquote>
<p>Also in my case, the time is not too different from just doing loss.item() every time.</p>
</blockquote>
</aside>
<p>This could mean that your code is already bottlenecks e.g. by other synchronizations.<br>
Profile the code using the PyTorch profiler or e.g. Nsight systems to see where the botleneck in the code is.</p>