()
| 81 | |
| 82 | |
| 83 | def main(): |
| 84 | import torch.multiprocessing |
| 85 | torch.multiprocessing.set_sharing_strategy('file_system') |
| 86 | |
| 87 | device = torch.device('cuda') |
| 88 | path = osp.join(osp.expanduser('~'), 'datasets', 'Reddit') |
| 89 | dataset = Reddit(path) |
| 90 | data = dataset[0].to(device) |
| 91 | |
| 92 | train_loader = NeighborSampler(data.edge_index, node_idx=None, |
| 93 | sizes=[10, 10, 25], batch_size=128, |
| 94 | shuffle=True, num_workers=32) |
| 95 | test_loader = NeighborSampler(data.edge_index, node_idx=None, |
| 96 | sizes=[10, 10, 25], batch_size=128, |
| 97 | shuffle=False, num_workers=32) |
| 98 | |
| 99 | gconv = GConv(input_dim=dataset.num_features, hidden_dim=512, num_layers=3).to(device) |
| 100 | encoder_model = Encoder(encoder=gconv, hidden_dim=512).to(device) |
| 101 | contrast_model = SingleBranchContrast(loss=L.JSD(), mode='G2L').to(device) |
| 102 | |
| 103 | optimizer = Adam(encoder_model.parameters(), lr=0.0001) |
| 104 | |
| 105 | with tqdm(total=30, desc='(T)') as pbar: |
| 106 | for epoch in range(1, 31): |
| 107 | loss = train(encoder_model, contrast_model, data, train_loader, optimizer) |
| 108 | pbar.set_postfix({'loss': loss}) |
| 109 | pbar.update() |
| 110 | |
| 111 | test_result = test(encoder_model, data, test_loader) |
| 112 | print(f'(E): Best test F1Mi={test_result["micro_f1"]:.4f}, F1Ma={test_result["macro_f1"]:.4f}') |
| 113 | |
| 114 | |
| 115 | if __name__ == '__main__': |
no test coverage detected