Fix dataset input bug & Logging graph indexing of patterns

.gitignore

@@ -5,7 +5,6 @@ log-*/
 *.p
 *.png
 plots/
-data/
 runs/
 results/
 *.sw?

analyze.py

@@ -0,0 +1,30 @@
+import sys
+
+count_pattern_by_size = {}
+pattern_list = {}
+pid = 0
+
+with open(sys.argv[1], 'r', encoding='utf-8') as file:
+    pattern_size = 0
+
+    for line in file:
+        if "Saving plots" in line:
+            plot_name = line.split("/")[-1]
+            pattern_size = int(plot_name.split("-")[0])
+
+            if pattern_size not in count_pattern_by_size:
+                count_pattern_by_size[pattern_size] = 1
+            else:
+                count_pattern_by_size[pattern_size] += 1
+
+        if "{" in line and "}" in line:
+            pattern_list[pid] = (pattern_size, [int(x) for x in line[1:-2].split(", ")])
+            pid += 1
+
+print("Count pattern")
+for k, v in count_pattern_by_size.items():
+    print(k, v)
+
+print("Pattern:")
+for k, v in pattern_list.items():
+    print(v[0])

common/data.py

@@ -47,6 +47,8 @@ def load_dataset(name):
         dataset = TUDataset(root="/tmp/FIRSTMM_DB", name="FIRSTMM_DB")
     elif name == "dblp":
         dataset = TUDataset(root="/tmp/DBLP_v1", name="DBLP_v1")
+    elif name == "msrc":
+        dataset = TUDataset(root="/tmp/MSRC-21", name="MSRC-21")
     elif name == "ppi":
         dataset = PPI(root="/tmp/PPI")
     elif name == "qm9":
@@ -426,6 +428,77 @@ def add_anchor(g):
         self.batch_idx += 1
         return pos_a, pos_b, neg_a, neg_b
 
+class  PreloadedDataSource(DataSource):
+    def __init__(self, dataset_name, tag=None):
+        self.dataset = dataset_name
+        self.train_keys = []
+        self.test_keys = []
+
+        with open (os.path.join(self.dataset, "train_keys.pkl"), 'rb') as fp:
+            self.train_keys = pickle.load(fp)
+            random.shuffle(self.train_keys)
+        if tag is None:
+            test_key_file = "test_keys.pkl"
+        else:
+            test_key_file = "_".join(["test_keys", tag]) + ".pkl"
+        with open (os.path.join(self.dataset, test_key_file), 'rb') as fp:
+            self.test_keys = pickle.load(fp)
+
+        self.train_size = len(self.train_keys)
+        self.test_size = len(self.test_keys)
+
+    def gen_data_loaders(self, _, batch_size, train=True, use_distributed_sampling=False):
+        if train:
+            size = self.train_size
+        else:
+            size = self.test_size
+        loop_time = size // batch_size
+        if size % batch_size != 0:
+            loop_time += 1
+
+        loaders = [[batch_size]*loop_time, [], []]
+        for i in range(loop_time):
+            if i == 0:
+                loaders[1].append(0)
+                loaders[2].append(min(size, batch_size))
+            else:
+                loaders[1].append(loaders[2][-1])
+                loaders[2].append(min(size, loaders[1][-1] + batch_size))
+
+        return loaders
+
+    def gen_batch(self, batch_size, start, end, train):
+        pos_a, pos_b, neg_a, neg_b = [], [], [], []
+        keys = []
+        if train:
+            keys = self.train_keys
+        else:
+            keys = self.test_keys
+
+        for key in keys[start:end]:
+            with open(os.path.join(self.dataset, key), 'rb') as f:
+                data = pickle.load(f)
+                if len(data) == 3:
+                    m1, m2, _ = data
+                else:
+                    m1, m2 = data
+
+                if "iso" in key:
+                    pos_a.append(m2)
+                    pos_b.append(m1)
+                else:
+                    neg_a.append(m2)
+                    neg_b.append(m1)
+
+        if pos_a: 
+            pos_a = utils.batch_nx_graphs(pos_a)
+            pos_b = utils.batch_nx_graphs(pos_b)
+        if neg_a:
+            neg_a = utils.batch_nx_graphs(neg_a)
+            neg_b = utils.batch_nx_graphs(neg_b)
+
+        return pos_a, pos_b, neg_a, neg_b
+
 if __name__ == "__main__":
     import matplotlib.pyplot as plt
     plt.rcParams.update({"font.size": 14})

common/models.py

@@ -230,6 +230,9 @@ def forward(self, x, edge_index, edge_weight=None, size=None,
         """
         #edge_index, edge_weight = add_remaining_self_loops(
         #    edge_index, edge_weight, 1, x.size(self.node_dim))
+        if edge_weight is None:
+            edge_weight = torch.ones((edge_index.size(1),), dtype=edge_index.dtype,
+                                 device=edge_index.device)
         edge_index, _ = pyg_utils.remove_self_loops(edge_index)
 
         return self.propagate(edge_index, size=size, x=x,

common/utils.py

@@ -15,6 +15,15 @@
 
 from common import feature_preprocess
 
+def set_seed(seed):
+    random.seed(seed)
+    # torch.backends.cudnn.deterministic=True
+    # torch.backends.cudnn.benchmark = False
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.cuda.manual_seed_all(seed)
+
 def sample_neigh(graphs, size):
     ps = np.array([len(g) for g in graphs], dtype=np.float)
     ps /= np.sum(ps)
@@ -58,7 +67,7 @@ def wl_hash(g, dim=64, node_anchored=False):
                 vecs[v] = 1
                 break
     for i in range(len(g)):
-        newvecs = np.zeros((len(g), dim), dtype=np.int)
+        newvecs = np.zeros((len(g), dim), dtype=np.int64)
         for n in g.nodes:
             newvecs[n] = vec_hash(np.sum(vecs[list(g.neighbors(n)) + [n]],
                 axis=0))
@@ -226,7 +235,7 @@ def batch_nx_graphs(graphs, anchors=None):
     #loader = DataLoader(motifs_batch, batch_size=len(motifs_batch))
     #for b in loader: batch = b
     augmenter = feature_preprocess.FeatureAugment()
-    
+
     if anchors is not None:
         for anchor, g in zip(anchors, graphs):
             for v in g.nodes: