| 521 | # ------------------------------------------------------------------------ |
| 522 | # 5) Evaluation function |
| 523 | # ------------------------------------------------------------------------ |
| 524 | def evaluate(model, val_dloader, device, distributed): |
| 525 | model.eval() |
| 526 | |
| 527 | total_loss = 0.0 |
| 528 | total_count = 0 |
| 529 | with torch.no_grad(): |
| 530 | for batch in val_dloader: |
| 531 | embeddings_1024 = batch["embeddings"].to(device) |
| 532 | texts = batch["texts"] |
| 533 | seq_lens = batch["seq_lens"] |
| 534 | |
| 535 | with autocast(enabled=False): |
| 536 | loss = model(embeddings_1024, texts, seq_lens) |
| 537 | |
| 538 | bs = embeddings_1024.size(0) |
| 539 | total_loss += loss.item() * bs |
| 540 | total_count += bs |
| 541 | |
| 542 | if distributed: |
| 543 | result = torch.tensor([total_loss, total_count], device=device, dtype=torch.float32) |
| 544 | dist.all_reduce(result, op=dist.ReduceOp.SUM) |
| 545 | total_loss, total_count = result[0].item(), result[1].item() |
| 546 | |
| 547 | model.train() |
| 548 | return (total_loss / total_count) if total_count > 0 else 0.0 |
| 549 | |
| 550 | |