Merge pull request #79 from joshchang1112:master

PiperOrigin-RevId: 351154923

Author: tensorflow-copybara

research/gnn-survey/README.md

@@ -3,16 +3,18 @@
 This repository contains a tensorflow 2.0 implementation of GNN models for node
 classification.
 
-Update: ** Sparse version of GCN and GAT are available **
+Update: ** Sparse version of GCN, GAT, GIN are available **
 
 ## Code organization
 
-*   `download_dataset.sh`: Download graph dataset. Now only `cora` is available.
+*   `download_dataset.sh`: Download graph dataset. Now `cora` and `citeseer` are
+    available.
 
 *   `train.py`: Trains a model with FLAGS parameters. `python train.py --help`
     for more information.
 
-*   `models.py`: Gnn models implementation. Now `gcn` and `gat` are available.
+*   `models.py`: Gnn models implementation. Now `gcn`, `gat` and `gin` are
+    available.
 
 *   `layers.py`: Single gnn layers implementation.
 
@@ -25,7 +27,7 @@ Update: ** Sparse version of GCN and GAT are available **
 2.  Download the dataset.
 
 ```
-bash download_dataset.sh
+bash download_dataset.sh <DATASET>
 ```
 
 1.  Train GAT on cora with default parameters.
@@ -39,7 +41,8 @@ python train.py --save_dir=$SAVE_DIR
 
 ## Training Results
 
-*   Better GAT results on cora (84.7% average test accuracy)[^GAT]:
+*   Better GAT results on cora (84.7% average test accuracy,
+    [[2]](#references)):
 
 ```
 python train.py \
@@ -56,7 +59,8 @@ python train.py \
   --sparse_features=True
 ```
 
-*   Reproduce gcn results on cora (81.5% average test accuracy)[^GCN]:
+*   Reproduce gcn results on cora (81.5% average test accuracy,
+    [[1]](#references)):
 
 ```
 python train.py \
@@ -72,7 +76,8 @@ python train.py \
   --sparse_features=True
 ```
 
-*   Better gcn results on cora (82.5% average test accuracy):
+*   Better gcn results on cora (82.5% average test accuracy,
+    [[1]](#references)):
 
 ```
 python train.py \
@@ -88,7 +93,7 @@ python train.py \
   --sparse_features=True
 ```
 
-*   GIN results on cora (81.7% average test accuracy):
+*   GIN results on cora (81.7% average test accuracy, [[3]](#references)):
 
 ```
 python train.py \
@@ -107,8 +112,11 @@ python train.py \
 
 ## References
 
-[^GCN]: Thomas N. Kipf, Max Welling "Semi-Supervised Classification with Graph
-    Convolutional Networks"
-    [GCN original github](https://github.com/tkipf/gcn/tree/master/gcn)
-[^GAT]: Petar Veličković, Guillem Cucurull, et al. "Graph Attention Networks"
-    [GAT original github](https://github.com/PetarV-/GAT)
+[[1] T. Kipf and M. Welling. "Semi-Supervised Classification with Graph
+Convolutional Networks" ICLR 2017](https://arxiv.org/pdf/1609.02907.pdf)
+
+[[2] P. Veličković, G. Cucurull, A. Casanova, A. Romero, P. Liò and Y. Bengio.
+"Graph Attention Networks" ICLR 2018](https://arxiv.org/pdf/1710.10903.pdf)
+
+[[3] K. Xu, W. Hu, J. Leskovec and S. Jegelka. "How Powerful are Graph Neural
+Networks?" ICLR 2019](https://arxiv.org/pdf/1810.00826.pdf)

research/gnn-survey/download_dataset.sh

@@ -13,8 +13,9 @@
 # limitations under the License.
 # TODO(ppham27): Consider consolidating with
 # examples/preprocess/cora/prep_data.sh.
-# URL for downloading Cora dataset.
-URL=https://linqs-data.soe.ucsc.edu/public/lbc/cora.tgz
+# URL for downloading dataset (Available dataset: cora, citeseer).
+DATASET=$1
+URL=https://linqs-data.soe.ucsc.edu/public/lbc/${DATASET}.tgz
 
 # Target folder to store and process data.
 DATA_DIR=data
@@ -32,6 +33,6 @@ function download () {
   fi
 }
 
-# Download and unzip the dataset. Data will be at '${DATA_DIR}/cora/' folder.
+# Download and unzip the dataset. Data will be at '${DATA_DIR}/${DATASET}/' folder.
 download ${URL} ${DATA_DIR}
-tar -C ${DATA_DIR} -xvzf ${DATA_DIR}/cora.tgz
+tar -C ${DATA_DIR} -xvzf ${DATA_DIR}/${DATASET}.tgz

research/gnn-survey/train.py

@@ -20,8 +20,8 @@
 
 from utils import load_dataset, build_model, cal_acc  # pylint: disable=g-multiple-import
 
-flags.DEFINE_enum('dataset', 'cora', ['cora'],
-                  'The input dataset. Avaliable dataset now: cora')
+flags.DEFINE_enum('dataset', 'cora', ['cora', 'citeseer'],
+                  'The input dataset. Avaliable dataset now: cora, citeseer')
 flags.DEFINE_enum('model', 'gat', ['gcn', 'gat', 'gin'],
                   'GNN model. Available model now: gcn, gat')
 flags.DEFINE_float('dropout_rate', 0.6, 'Dropout probability')

research/gnn-survey/utils.py

@@ -65,24 +65,34 @@ def cal_acc(labels, logits):
   return acc.numpy().item()
 
 
-def encode_onehot(labels):
+def encode_onehot(dataset, labels):
   """Provides a mapping from string labels to integer indices."""
   label_index = {
-      'Case_Based': 0,
-      'Genetic_Algorithms': 1,
-      'Neural_Networks': 2,
-      'Probabilistic_Methods': 3,
-      'Reinforcement_Learning': 4,
-      'Rule_Learning': 5,
-      'Theory': 6,
+      'cora': {
+          'Case_Based': 0,
+          'Genetic_Algorithms': 1,
+          'Neural_Networks': 2,
+          'Probabilistic_Methods': 3,
+          'Reinforcement_Learning': 4,
+          'Rule_Learning': 5,
+          'Theory': 6
+      },
+      'citeseer': {
+          'AI': 0,
+          'IR': 1,
+          'HCI': 2,
+          'DB': 3,
+          'ML': 4,
+          'Agents': 5
+      }
   }
 
   # Convert to onehot label
-  num_classes = len(label_index)
+  num_classes = len(label_index[dataset])
   onehot_labels = np.zeros((len(labels), num_classes))
   idx = 0
   for s in labels:
-    onehot_labels[idx, label_index[s]] = 1
+    onehot_labels[idx, label_index[dataset][s]] = 1
     idx += 1
   return onehot_labels
 
@@ -115,23 +125,30 @@ def sparse_matrix_to_tf_sparse_tensor(matrix):
 
 
 def load_dataset(dataset, sparse_features, normalize_adj):
-  """Loads Cora dataset."""
+  """Loads dataset."""
   dir_path = os.path.join('data', dataset)
   content_path = os.path.join(dir_path, '{}.content'.format(dataset))
   citation_path = os.path.join(dir_path, '{}.cites'.format(dataset))
 
   content = np.genfromtxt(content_path, dtype=np.dtype(str))
-
-  idx = np.array(content[:, 0], dtype=np.int32)
+  idx = np.array(content[:, 0])
   features = sp.csr_matrix(content[:, 1:-1], dtype=np.float32)
-  labels = encode_onehot(content[:, -1])
+  labels = encode_onehot(dataset, content[:, -1])
 
   # Dict which maps paper id to data id
   idx_map = {j: i for i, j in enumerate(idx)}
-  edges_unordered = np.genfromtxt(citation_path, dtype=np.int32)
+  edges_unordered = np.genfromtxt(citation_path, dtype=np.dtype(str))
   edges = np.array(
       list(map(idx_map.get, edges_unordered.flatten())),
-      dtype=np.int32).reshape(edges_unordered.shape)
+      dtype=np.dtype(str)).reshape(edges_unordered.shape)
+
+  # Delete relation which the nodes appear in cites but not in content
+  del_rel = []
+  for i, j in enumerate(edges):
+    if j[0] == 'None' or j[1] == 'None':
+      del_rel.append(i)
+  edges = np.delete(edges, del_rel, 0)
+
   adj = sp.coo_matrix((np.ones(edges.shape[0]), (edges[:, 0], edges[:, 1])),
                       shape=(labels.shape[0], labels.shape[0]),
                       dtype=np.float32)
@@ -145,13 +162,13 @@ def load_dataset(dataset, sparse_features, normalize_adj):
   if normalize_adj:
     adj = normalize_adj_matrix(adj)
 
-  # 5% for train, 300 for validation, 1000 for test
   idx_train = slice(140)
   idx_val = slice(200, 500)
   idx_test = slice(500, 1500)
 
   features = tf.convert_to_tensor(np.array(features.todense()))
   labels = tf.convert_to_tensor(np.where(labels)[1])
+
   if sparse_features:
     adj = sparse_matrix_to_tf_sparse_tensor(adj)
   else: