| 52 | |
| 53 | |
| 54 | def train(encoder_model, contrast_model, data, dataloader, optimizer): |
| 55 | encoder_model.train() |
| 56 | total_loss = total_examples = 0 |
| 57 | for batch_size, node_id, adjs in dataloader: |
| 58 | adjs = [adj.to('cuda') for adj in adjs] |
| 59 | optimizer.zero_grad() |
| 60 | z, g, zn = encoder_model(data.x[node_id], adjs) |
| 61 | loss = contrast_model(h=z, g=g, hn=zn) |
| 62 | loss.backward() |
| 63 | optimizer.step() |
| 64 | total_loss += loss.item() * z.shape[0] |
| 65 | total_examples += z.shape[0] |
| 66 | return total_loss / total_examples |
| 67 | |
| 68 | |
| 69 | def test(encoder_model, data, dataloader): |