()
| 89 | |
| 90 | |
| 91 | def main(): |
| 92 | device = torch.device('cuda') |
| 93 | path = osp.join(osp.expanduser('~'), 'datasets') |
| 94 | dataset = Planetoid(path, name='Cora', transform=T.NormalizeFeatures()) |
| 95 | data = dataset[0].to(device) |
| 96 | |
| 97 | aug1 = A.Compose([A.EdgeRemoving(pe=0.3), A.FeatureMasking(pf=0.3)]) |
| 98 | aug2 = A.Compose([A.EdgeRemoving(pe=0.3), A.FeatureMasking(pf=0.3)]) |
| 99 | |
| 100 | gconv = GConv(input_dim=dataset.num_features, hidden_dim=32, activation=torch.nn.ReLU, num_layers=2).to(device) |
| 101 | encoder_model = Encoder(encoder=gconv, augmentor=(aug1, aug2), hidden_dim=32, proj_dim=32).to(device) |
| 102 | contrast_model = DualBranchContrast(loss=L.InfoNCE(tau=0.2), mode='L2L', intraview_negs=True).to(device) |
| 103 | |
| 104 | optimizer = Adam(encoder_model.parameters(), lr=0.01) |
| 105 | |
| 106 | with tqdm(total=1000, desc='(T)') as pbar: |
| 107 | for epoch in range(1, 1001): |
| 108 | loss = train(encoder_model, contrast_model, data, optimizer) |
| 109 | pbar.set_postfix({'loss': loss}) |
| 110 | pbar.update() |
| 111 | |
| 112 | test_result = test(encoder_model, data) |
| 113 | print(f'(E): Best test F1Mi={test_result["micro_f1"]:.4f}, F1Ma={test_result["macro_f1"]:.4f}') |
| 114 | |
| 115 | |
| 116 | if __name__ == '__main__': |
no test coverage detected