| 254 | ################################# |
| 255 | |
| 256 | def train_iter(epoch, model, optimizer, criterion, input, target, mask_train, mask_val, print_every=10): |
| 257 | start_t = time.time() |
| 258 | model.train() |
| 259 | optimizer.zero_grad() |
| 260 | |
| 261 | # Forward pass |
| 262 | output = model(*input) |
| 263 | loss = criterion(output[mask_train], target[mask_train]) # Compute the loss using the training mask |
| 264 | |
| 265 | loss.backward() |
| 266 | optimizer.step() |
| 267 | |
| 268 | # Evaluate the model performance on training and validation sets |
| 269 | loss_train, acc_train = test(model, criterion, input, target, mask_train) |
| 270 | loss_val, acc_val = test(model, criterion, input, target, mask_val) |
| 271 | |
| 272 | if epoch % print_every == 0: |
| 273 | # Print the training progress at specified intervals |
| 274 | print(f'Epoch: {epoch:04d} ({(time.time() - start_t):.4f}s) loss_train: {loss_train:.4f} acc_train: {acc_train:.4f} loss_val: {loss_val:.4f} acc_val: {acc_val:.4f}') |
| 275 | |
| 276 | |
| 277 | def test(model, criterion, input, target, mask): |