| 67 | |
| 68 | |
| 69 | def train(encoder_model, contrast_model, dataloader, optimizer): |
| 70 | encoder_model.train() |
| 71 | epoch_loss = 0 |
| 72 | for data in dataloader: |
| 73 | data = data.to('cuda') |
| 74 | optimizer.zero_grad() |
| 75 | |
| 76 | if data.x is None: |
| 77 | num_nodes = data.batch.size(0) |
| 78 | data.x = torch.ones((num_nodes, 1), dtype=torch.float32, device=data.batch.device) |
| 79 | |
| 80 | _, _, _, _, g1, g2 = encoder_model(data.x, data.edge_index, data.batch) |
| 81 | g1, g2 = [encoder_model.encoder.project(g) for g in [g1, g2]] |
| 82 | loss = contrast_model(g1=g1, g2=g2, batch=data.batch) |
| 83 | loss.backward() |
| 84 | optimizer.step() |
| 85 | |
| 86 | epoch_loss += loss.item() |
| 87 | return epoch_loss |
| 88 | |
| 89 | |
| 90 | def test(encoder_model, dataloader): |