| 33 | |
| 34 | |
| 35 | def get_split(dataset, y, nclass, seed=0): |
| 36 | |
| 37 | if dataset == 'arxiv': |
| 38 | dataset = DglNodePropPredDataset('ogbn-arxiv') |
| 39 | split = dataset.get_idx_split() |
| 40 | train, valid, test = split['train'], split['valid'], split['test'] |
| 41 | return train, valid, test |
| 42 | |
| 43 | elif dataset == 'penn': |
| 44 | split = np.load('node_raw_data/fb100-Penn94-splits.npy', allow_pickle=True)[0] |
| 45 | train, valid, test = split['train'], split['valid'], split['test'] |
| 46 | return train, valid, test |
| 47 | |
| 48 | else: |
| 49 | y = y.cpu() |
| 50 | |
| 51 | percls_trn = int(round(0.6 * len(y) / nclass)) |
| 52 | val_lb = int(round(0.2 * len(y))) |
| 53 | |
| 54 | indices = [] |
| 55 | for i in range(nclass): |
| 56 | index = (y == i).nonzero().view(-1) |
| 57 | index = index[torch.randperm(index.size(0), device=index.device)] |
| 58 | indices.append(index) |
| 59 | |
| 60 | train_index = torch.cat([i[:percls_trn] for i in indices], dim=0) |
| 61 | rest_index = torch.cat([i[percls_trn:] for i in indices], dim=0) |
| 62 | rest_index = rest_index[torch.randperm(rest_index.size(0))] |
| 63 | valid_index = rest_index[:val_lb] |
| 64 | test_index = rest_index[val_lb:] |
| 65 | |
| 66 | return train_index, valid_index, test_index |
| 67 | |