(dataloader, model, loss_fn, optimizer)
| 136 | # backpropagates the prediction error to adjust the model's parameters. |
| 137 | |
| 138 | def train(dataloader, model, loss_fn, optimizer): |
| 139 | size = len(dataloader.dataset) |
| 140 | model.train() |
| 141 | for batch, (X, y) in enumerate(dataloader): |
| 142 | X, y = X.to(device), y.to(device) |
| 143 | |
| 144 | # Compute prediction error |
| 145 | pred = model(X) |
| 146 | loss = loss_fn(pred, y) |
| 147 | |
| 148 | # Backpropagation |
| 149 | loss.backward() |
| 150 | optimizer.step() |
| 151 | optimizer.zero_grad() |
| 152 | |
| 153 | if batch % 100 == 0: |
| 154 | loss, current = loss.item(), (batch + 1) * len(X) |
| 155 | print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]") |
| 156 | |
| 157 | ############################################################################## |
| 158 | # We also check the model's performance against the test dataset to ensure it is learning. |
no test coverage detected