| 28 | |
| 29 | @torch.no_grad() |
| 30 | def test(model, test_loader, dataset_name): |
| 31 | evaluator = Evaluator(name=dataset_name) |
| 32 | model.eval() |
| 33 | xs = [] |
| 34 | y_true = [] |
| 35 | for i, batch in enumerate(test_loader): |
| 36 | if i == 0: |
| 37 | device = batch.x.device |
| 38 | batch.x = batch.x.to(torch.float32) # TODO |
| 39 | x = model(batch.x, batch.edge_index)[: batch.batch_size] |
| 40 | xs.append(x.cpu()) |
| 41 | y_true.append(batch.y[: batch.batch_size].cpu()) |
| 42 | del batch |
| 43 | |
| 44 | xs = [t.to(device) for t in xs] |
| 45 | y_true = [t.to(device) for t in y_true] |
| 46 | y_pred = torch.cat(xs, dim=0).argmax(dim=-1, keepdim=True) |
| 47 | y_true = torch.cat(y_true, dim=0).unsqueeze(-1) |
| 48 | test_acc = evaluator.eval( |
| 49 | { |
| 50 | "y_true": y_true, |
| 51 | "y_pred": y_pred, |
| 52 | } |
| 53 | )["acc"] |
| 54 | return test_acc |
| 55 | |
| 56 | |
| 57 | gs.set_option(show_log=True) |