update amazon

Author: UNAOUN

examples/glnn/train_student.py

@@ -12,7 +12,7 @@
 import tensorlayerx as tlx
 from gammagl.datasets import Planetoid, Amazon
 from gammagl.models import MLP
-from gammagl.utils import mask_to_index, get_train_val_test_split
+from gammagl.utils import mask_to_index
 from tensorlayerx.model import TrainOneStep, WithLoss
 
 
@@ -73,25 +73,18 @@ def train_student(args):
     if args.dataset in ['cora', 'pubmed', 'citeseer']:
         dataset = Planetoid(args.dataset_path, args.dataset)
     elif args.dataset in ['computers', 'photo']:
-        dataset = Amazon(args.dataset_path, args.dataset)
+        dataset = Amazon(args.dataset_path, args.dataset, train_per_class=20, val_per_class=30)
     graph = dataset[0]
 
     # load teacher_logits from .npy file
     teacher_logits = tlx.files.load_npy_to_any(path = r'./', name = f'{args.dataset}_{args.teacher}_logits.npy')
     teacher_logits = tlx.ops.convert_to_tensor(teacher_logits)
 
     # for mindspore, it should be passed into node indices
-    if args.dataset in ['cora', 'pubmed', 'citeseer']:
-        train_idx = mask_to_index(graph.train_mask)
-        test_idx = mask_to_index(graph.test_mask)
-        val_idx = mask_to_index(graph.val_mask)
-        t_idx = tlx.concat([train_idx, test_idx, val_idx], axis=0)
-    elif args.dataset in ['computers', 'photo']:
-        train_mask, val_mask, test_mask = get_train_val_test_split(dataset, train_per_class=20, val_per_class=30)
-        train_idx = mask_to_index(train_mask)
-        val_idx = mask_to_index(val_mask)
-        test_idx = mask_to_index(test_mask)
-        t_idx = tlx.concat([train_idx, test_idx, val_idx], axis=0)
+    train_idx = mask_to_index(graph.train_mask)
+    test_idx = mask_to_index(graph.test_mask)
+    val_idx = mask_to_index(graph.val_mask)
+    t_idx = tlx.concat([train_idx, test_idx, val_idx], axis=0)
 
     net = MLP(in_channels=dataset.num_node_features,
               hidden_channels=conf["hidden_dim"],

examples/glnn/train_teacher.py

@@ -12,7 +12,7 @@
 import tensorlayerx as tlx
 from gammagl.datasets import Planetoid, Amazon
 from gammagl.models import GCNModel, GraphSAGE_Full_Model, GATModel, APPNPModel, MLP
-from gammagl.utils import mask_to_index, calc_gcn_norm, get_train_val_test_split
+from gammagl.utils import mask_to_index, calc_gcn_norm
 from tensorlayerx.model import TrainOneStep, WithLoss
 
 
@@ -66,21 +66,15 @@ def train_teacher(args):
     if args.dataset in ['cora', 'pubmed', 'citeseer']:
         dataset = Planetoid(args.dataset_path, args.dataset)
     elif args.dataset in ['computers', 'photo']:
-        dataset = Amazon(args.dataset_path, args.dataset)
+        dataset = Amazon(args.dataset_path, args.dataset, train_per_class=20, val_per_class=30)
     graph = dataset[0]
     edge_index = graph.edge_index
     edge_weight = tlx.convert_to_tensor(calc_gcn_norm(edge_index, graph.num_nodes))
 
     # for mindspore, it should be passed into node indices
-    if args.dataset in ['cora', 'pubmed', 'citeseer']:
-        train_idx = mask_to_index(graph.train_mask)
-        test_idx = mask_to_index(graph.test_mask)
-        val_idx = mask_to_index(graph.val_mask)
-    elif args.dataset in ['computers', 'photo']:
-        train_mask, val_mask, test_mask = get_train_val_test_split(dataset, train_per_class=20, val_per_class=30)
-        train_idx = mask_to_index(train_mask)
-        val_idx = mask_to_index(val_mask)
-        test_idx = mask_to_index(test_mask)
+    train_idx = mask_to_index(graph.train_mask)
+    test_idx = mask_to_index(graph.test_mask)
+    val_idx = mask_to_index(graph.val_mask)
 
     if args.teacher == "GCN":
         net = GCNModel(feature_dim=dataset.num_node_features, 

gammagl/datasets/amazon.py

@@ -3,6 +3,7 @@
 import tensorlayerx as tlx
 from gammagl.data import InMemoryDataset, download_url
 from gammagl.io.npz import read_npz
+from gammagl.utils import get_train_val_test_split
 
 
 class Amazon(InMemoryDataset):
@@ -33,6 +34,10 @@ class Amazon(InMemoryDataset):
         being saved to disk. (default: :obj:`None`)
     force_reload (bool, optional): Whether to re-process the dataset.
         (default: :obj:`False`)
+    train_per_class (int, optional): Number of training samples per class.
+        (default: :obj:`20`)
+    val_per_class (int, optional): Number of validation samples per class.
+        (default: :obj:`20`)
 
     Stats:
         .. list-table::
@@ -61,12 +66,18 @@ class Amazon(InMemoryDataset):
     def __init__(self, root: str = None, name: str = 'computers',
                  transform: Optional[Callable] = None,
                  pre_transform: Optional[Callable] = None,
-                 force_reload: bool = False):
+                 force_reload: bool = False,
+                 train_per_class: int = 20,
+                 val_per_class: int = 20):
         self.name = name.lower()
         assert self.name in ['computers', 'photo']
         super().__init__(root, transform, pre_transform, force_reload = force_reload)
         self.data, self.slices = self.load_data(self.processed_paths[0])
 
+        data = self.get(0)
+        data.train_mask, data.val_mask, data.test_mask = get_train_val_test_split(self.data, train_per_class, val_per_class, self.num_classes)
+        self.data, self.slices = self.collate([data])
+
     @property
     def raw_dir(self) -> str:
         return osp.join(self.root, self.name.capitalize(), 'raw')

gammagl/utils/get_split.py

@@ -2,27 +2,28 @@
 import numpy as np
 
 
-def get_train_val_test_split(dataset, train_per_class, val_per_class):
+def get_train_val_test_split(graph, train_per_class, val_per_class, num_classes):
     """Split the dataset into train, validation, and test sets.
 
     Parameters
     ----------
-    dataset :
-        The dataset to split.
+    graph :
+        The graph to split.
     train_per_class : int
         The number of training examples per class.
     val_per_class : int
         The number of validation examples per class.
+    num_classes : int
+        The number of classes in the dataset.
 
     Returns
     -------
     :class:`tuple` of :class:`tensor`
     
     """
-    graph = dataset[0]
     random_state = np.random.RandomState(0)
-    labels = tlx.nn.OneHot(depth=dataset.num_classes)(graph.y).numpy()
-    num_samples, num_classes = graph.num_nodes, dataset.num_classes
+    labels = tlx.nn.OneHot(depth=num_classes)(graph.y).numpy()
+    num_samples, num_classes = graph.num_nodes, num_classes
     remaining_indices = set(range(num_samples))
     forbidden_indices = set()