use tu dataset

pyg-team · Jul 6, 2019 · 9a07e84 · 9a07e84
1 parent e00adba
commit 9a07e84
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 36 deletions.
diff --git a/examples/colors_topk_pool.py b/examples/colors_topk_pool.py
@@ -4,7 +4,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch_geometric.datasets import SyntheticDataset
+from torch_geometric.datasets import TUDataset
 from torch_geometric.transforms import HandleNodeAttention
 from torch_geometric.data import DataLoader
 from torch_geometric.nn import GINConv, TopKPooling
@@ -13,8 +13,8 @@
 
 train_path = osp.join(
     osp.dirname(osp.realpath(__file__)), '..', 'data', 'COLORS-3')
-dataset = SyntheticDataset(train_path, name='COLORS-3', use_node_attr=True,
-                           transform=HandleNodeAttention())
+dataset = TUDataset(train_path, name='COLORS-3', use_node_attr=True,
+                    transform=HandleNodeAttention())
 
 n_train, n_val, n_test_each = 500, 2500, 2500
 

diff --git a/examples/triangles_sag_pool.py b/examples/triangles_sag_pool.py
@@ -4,7 +4,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch_geometric.datasets import SyntheticDataset
+from torch_geometric.datasets import TUDataset
 from torch_geometric.transforms import OneHotDegree, HandleNodeAttention
 from torch_geometric.transforms import Compose
 from torch_geometric.data import DataLoader
@@ -16,8 +16,8 @@
 
 train_path = osp.join(
     osp.dirname(osp.realpath(__file__)), '..', 'data', 'TRIANGLES')
-dataset = SyntheticDataset(train_path, name='TRIANGLES', use_node_attr=True,
-                           transform=transform)
+dataset = TUDataset(train_path, name='TRIANGLES', use_node_attr=True,
+                    transform=transform)
 
 n_train, n_val, n_test_each = 30000, 5000, 5000
 

diff --git a/torch_geometric/datasets/__init__.py b/torch_geometric/datasets/__init__.py
@@ -24,7 +24,6 @@
 from .bitcoin_otc import BitcoinOTC
 from .icews import ICEWS18
 from .gdelt import GDELT
-from .synthetic_dataset import SyntheticDataset
 
 __all__ = [
     'KarateClub',
@@ -53,5 +52,4 @@
     'BitcoinOTC',
     'ICEWS18',
     'GDELT',
-    'SyntheticDataset',
 ]
diff --git a/torch_geometric/datasets/synthetic_dataset.py b/torch_geometric/datasets/synthetic_dataset.py