dmlc / dmlc/dgl

[Examples] Update examples to record total loss via `total_loss += loss.detach()`

Open
#5,302 3 comments 0 reactions 0 assignees View on GitHub
stale-issue
Dominant language
Python
Stars
14.3k
Forks
3.1k
PR merge metrics
No merged PRs in 30d

Description

## 🚀 Feature
Most of the current examples have a training loops like one of the following:

Using `loss.item()`:
```python3
for x, y in dataloader:
y_pred = model(x)
loss = loss_fn(y, y_pred)
opt.zero_grad()
loss.backward()
opt.step()
total_loss += loss.item()
```
Calling `loss.item()` requires synchronizing the CPU and GPU, and prevents overlapping forward/backward/optimizer with the dataloader. See the Nsight Systems profile below:
![total_loss](https://user-images.githubusercontent.com/63612878/219457408-ac639c08-9dd5-4dfc-bb46-34efd491b2b3.png)

Using `loss` directly will avoid this:
```python3
for x, y in dataloader:
y_pred = model(x)
loss = loss_fn(y, y_pred)
opt.zero_grad()
loss.backward()
opt.step()
total_loss += loss
```
However, while in many cases this will not negatively impact the computation, it can end up connecting `total_loss` to the computation graph.

What we should be doing is calling `loss.detach()`:
```python3
for x, y in dataloader:
y_pred = model(x)
loss = loss_fn(y, y_pred)
opt.zero_grad()
loss.backward()
opt.step()
total_loss += loss.detach()
```

This accumulates the loss as a 1-element tensor on the training device, not attached to the computation graph, and allows the CPU to move on to the dataloader when the GPU is still working on forward/backward/optimizer.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.