| 545 | # 5) Evaluation function |
| 546 | # ------------------------------------------------------------------------ |
| 547 | def evaluate(model, val_dloader, device, distributed): |
| 548 | model.eval() |
| 549 | |
| 550 | total_loss = 0.0 |
| 551 | total_count = 0 |
| 552 | with torch.no_grad(): |
| 553 | for batch in val_dloader: |
| 554 | embeddings_1024 = batch["embeddings"].to(device) |
| 555 | texts = batch["texts"] |
| 556 | seq_lens = batch["seq_lens"] |
| 557 | |
| 558 | with autocast(enabled=False): |
| 559 | loss = model(embeddings_1024, texts, seq_lens) |
| 560 | |
| 561 | bs = embeddings_1024.size(0) |
| 562 | total_loss += loss.item() * bs |
| 563 | total_count += bs |
| 564 | |
| 565 | if distributed: |
| 566 | result = torch.tensor([total_loss, total_count], device=device, dtype=torch.float32) |
| 567 | dist.all_reduce(result, op=dist.ReduceOp.SUM) |
| 568 | total_loss, total_count = result[0].item(), result[1].item() |
| 569 | |
| 570 | model.train() |
| 571 | return (total_loss / total_count) if total_count > 0 else 0.0 |
| 572 | |
| 573 | # ------------------------------------------------------------------------ |
| 574 | # 6) Full-checkpoint Save & Load |