split the dataset into training set, validation set and testing set
(self, data, seed=717, train_size=0.7, val_size=0.1)
| 165 | self.save([data], self.processed_paths[0]) |
| 166 | |
| 167 | def _random_split(self, data, seed=717, train_size=0.7, val_size=0.1): |
| 168 | """split the dataset into training set, validation set and testing set""" |
| 169 | |
| 170 | assert 0 <= train_size + val_size <= 1, ( |
| 171 | "The sum of valid training set size and validation set size " |
| 172 | "must between 0 and 1 (inclusive)." |
| 173 | ) |
| 174 | |
| 175 | N = data.x.shape[0] |
| 176 | index = np.arange(N) |
| 177 | if self.name == "amazon": |
| 178 | # 0-3304 are unlabeled nodes |
| 179 | index = np.arange(3305, N) |
| 180 | |
| 181 | index = np.random.RandomState(seed).permutation(index) |
| 182 | train_idx = index[: int(train_size * len(index))] |
| 183 | val_idx = index[len(index) - int(val_size * len(index)):] |
| 184 | test_idx = index[ |
| 185 | int(train_size * len(index)): len(index) |
| 186 | - int(val_size * len(index)) |
| 187 | ] |
| 188 | train_mask = np.zeros(N, dtype=np.bool_) |
| 189 | val_mask = np.zeros(N, dtype=np.bool_) |
| 190 | test_mask = np.zeros(N, dtype=np.bool_) |
| 191 | train_mask[train_idx] = True |
| 192 | val_mask[val_idx] = True |
| 193 | test_mask[test_idx] = True |
| 194 | data.train_mask = torch.tensor(train_mask) |
| 195 | data.val_mask = torch.tensor(val_mask) |
| 196 | data.test_mask = torch.tensor(test_mask) |
| 197 | |
| 198 | return data |
| 199 | |
| 200 | def __repr__(self): |
| 201 | return f'{self.name}()' |