| 64 | |
| 65 | |
| 66 | def main(): |
| 67 | device = torch.device('cuda') |
| 68 | path = osp.join(osp.expanduser('~'), 'datasets', 'WikiCS') |
| 69 | dataset = WikiCS(path, transform=T.NormalizeFeatures()) |
| 70 | data = dataset[0].to(device) |
| 71 | |
| 72 | aug1 = A.Compose([A.EdgeRemoving(pe=0.5), A.FeatureMasking(pf=0.1)]) |
| 73 | aug2 = A.Compose([A.EdgeRemoving(pe=0.5), A.FeatureMasking(pf=0.1)]) |
| 74 | |
| 75 | gconv = GConv(input_dim=dataset.num_features, hidden_dim=256).to(device) |
| 76 | encoder_model = Encoder(encoder=gconv, augmentor=(aug1, aug2)).to(device) |
| 77 | contrast_model = WithinEmbedContrast(loss=L.BarlowTwins()).to(device) |
| 78 | |
| 79 | optimizer = Adam(encoder_model.parameters(), lr=5e-4) |
| 80 | scheduler = LinearWarmupCosineAnnealingLR( |
| 81 | optimizer=optimizer, |
| 82 | warmup_epochs=400, |
| 83 | max_epochs=4000) |
| 84 | |
| 85 | with tqdm(total=4000, desc='(T)') as pbar: |
| 86 | for epoch in range(1, 4001): |
| 87 | loss = train(encoder_model, contrast_model, data, optimizer) |
| 88 | scheduler.step() |
| 89 | pbar.set_postfix({'loss': loss}) |
| 90 | pbar.update() |
| 91 | |
| 92 | test_result = test(encoder_model, data) |
| 93 | print(f'(E): Best test F1Mi={test_result["micro_f1"]:.4f}, F1Ma={test_result["macro_f1"]:.4f}') |
| 94 | |
| 95 | |
| 96 | if __name__ == '__main__': |