Merge pull request #82 from otiliastr:fix_edges

PiperOrigin-RevId: 368331006

Author: tensorflow-copybara

research/gam/gam/data/dataset.py

@@ -253,12 +253,19 @@ class GraphDataset(Dataset):
   """Data container for SSL datasets."""
 
   class Edge(object):
+    """Graph Edge."""
 
     def __init__(self, src, tgt, weight=None):
       self.src = src
       self.tgt = tgt
       self.weight = weight
 
+    def copy(self, src=None, tgt=None, weight=None):
+      src = src if src is not None else self.src
+      tgt = tgt if tgt is not None else self.tgt
+      weight = weight if weight is not None else self.weight
+      return GraphDataset.Edge(src, tgt, weight)
+
   def __init__(self,
                name,
                features,
@@ -879,9 +886,8 @@ def restore_state_from_file(self, path):
     indices_type = (
         np.uint32
         if self.dataset.num_samples < np.iinfo(np.uint32).max else np.uint64)
-    if os.path.exists(file_indices_train) \
-        and os.path.exists(file_indices_unlabeled) \
-        and os.path.exists(file_labels_train):
+    if os.path.exists(file_indices_train) and os.path.exists(
+        file_indices_unlabeled) and os.path.exists(file_labels_train):
       with open(file_indices_train, 'r') as f:
         indices_train = np.genfromtxt(f, delimiter=',', dtype=indices_type)
       with open(file_indices_unlabeled, 'r') as f:

research/gam/gam/trainer/trainer_classification.py

@@ -661,9 +661,11 @@ def edge_iterator(self, data, batch_size, labeling):
       edges = data.get_edges(
           src_labeled=True, tgt_labeled=True, label_must_match=True)
     elif labeling == 'lu':
-      edges = (
-          data.get_edges(src_labeled=True, tgt_labeled=False) +
-          data.get_edges(src_labeled=False, tgt_labeled=True))
+      edges_lu = data.get_edges(src_labeled=True, tgt_labeled=False)
+      edges_ul = data.get_edges(src_labeled=False, tgt_labeled=True)
+      # Reverse the edges of UL to be LU.
+      edges_ul = [e.copy(src=e.tgt, tgt=e.src) for e in edges_ul]
+      edges = edges_lu + edges_ul
     elif labeling == 'uu':
       edges = data.get_edges(src_labeled=False, tgt_labeled=False)
     else:

research/gam/gam/trainer/trainer_classification_gcn.py

@@ -689,9 +689,11 @@ def edge_iterator(self, data, batch_size, labeling):
       edges = data.get_edges(
           src_labeled=True, tgt_labeled=True, label_must_match=True)
     elif labeling == 'lu':
-      edges = (
-          data.get_edges(src_labeled=True, tgt_labeled=False) +
-          data.get_edges(src_labeled=False, tgt_labeled=True))
+      edges_lu = data.get_edges(src_labeled=True, tgt_labeled=False)
+      edges_ul = data.get_edges(src_labeled=False, tgt_labeled=True)
+      # Reverse the edges of UL to be LU.
+      edges_ul = [e.copy(src=e.tgt, tgt=e.src) for e in edges_ul]
+      edges = edges_lu + edges_ul
     elif labeling == 'uu':
       edges = data.get_edges(src_labeled=False, tgt_labeled=False)
     else: