update the flex data code

graph_dit/datasets/dataset.py

@@ -70,7 +70,7 @@ class DataModule(AbstractDataModule):
         #     base_path = pathlib.Path(os.path.realpath(__file__)).parents[2]
         # except NameError:
         # base_path = pathlib.Path(os.getcwd()).parent[2]
-        base_path = '/home/stud/hanzhang/nasbenchDiT'
+        base_path = '/nfs/data3/hanzhang/nasbenchDiT'
         root_path = os.path.join(base_path, self.datadir)
         self.root_path = root_path
 
@@ -408,6 +408,7 @@ def new_graphs_to_json(graphs, filename):
         adj = graph[0]
 
         n_node = len(ops)
+        print(n_node)
         n_edge = len(ops)
         n_node_list.append(n_node)
         n_edge_list.append(n_edge)
@@ -489,7 +490,7 @@ def new_graphs_to_json(graphs, filename):
         'transition_E': transition_E.tolist(),
     }
 
-    with open(f'/home/stud/hanzhang/nasbenchDiT/graph_dit/nasbench-201-meta.json', 'w') as f:
+    with open(f'/nfs/data3/hanzhang/nasbenchDiT/graph_dit/nasbench-201-meta.json', 'w') as f:
         json.dump(meta_dict, f)
     
     return meta_dict
@@ -655,7 +656,7 @@ def graphs_to_json(graphs, filename):
 class Dataset(InMemoryDataset):
     def __init__(self, source, root, target_prop=None, transform=None, pre_transform=None, pre_filter=None):
         self.target_prop = target_prop
-        source = '/home/stud/hanzhang/nasbenchDiT/graph_dit/NAS-Bench-201-v1_1-096897.pth'
+        source = '/nfs/data3/hanzhang/nasbenchDiT/graph_dit/NAS-Bench-201-v1_1-096897.pth'
         self.source = source
         # self.api = API(source)  # Initialize NAS-Bench-201 API
         # print('API loaded')
@@ -676,7 +677,7 @@ class Dataset(InMemoryDataset):
         return [f'{self.source}.pt']
 
     def process(self):
-        source = '/home/stud/hanzhang/nasbenchDiT/graph_dit/NAS-Bench-201-v1_1-096897.pth'
+        source = '/nfs/data3/hanzhang/nasbenchDiT/graph_dit/NAS-Bench-201-v1_1-096897.pth'
         # self.api = API(source)
 
         data_list = []
@@ -712,6 +713,7 @@ class Dataset(InMemoryDataset):
         def generate_flex_adj_mat(ori_nodes, ori_edges, max_nodes=12, min_nodes=8,random_ratio=0.5):
             # print(ori_nodes)
             # print(ori_edges)
+            
             ori_edges = np.array(ori_edges)
             # ori_nodes = np.array(ori_nodes)
             nasbench_201_node_num = 8
@@ -720,8 +722,13 @@ class Dataset(InMemoryDataset):
             # print(f'arch_str: {arch_str}, \nmax_nodes: {max_nodes}, min_nodes: {min_nodes}, nodes_num: {nodes_num},random_seed: {random_seed},random_ratio: {random_ratio}')
             add_num = nodes_num - nasbench_201_node_num
             # ori_nodes, ori_edges = parse_architecture_string(arch_str)
-            add_nodes = [op for op in random.choices(num_to_op[1:-1], k=add_num)]
+            add_nodes = []
+            print(f'add_num: {add_num}')
+            for i in range(add_num):
+                add_nodes.append(random.choice(num_to_op[1:-1]))
             # print(add_nodes)
+            print(f'ori_nodes[:-1]: {ori_nodes[:-1]}, add_nodes: {add_nodes}')
+            print(f'len(ori_nodes[:-1]): {len(ori_nodes[:-1])}, len(add_nodes): {len(add_nodes)}')
             nodes = ori_nodes[:-1] + add_nodes + ['output']
             edges = np.zeros((nodes_num , nodes_num))
             edges[:6, :6] = ori_edges[:6, :6]
@@ -731,6 +738,11 @@ class Dataset(InMemoryDataset):
                     rand = random.random()
                     if rand < random_ratio:
                         edges[i, j] = 1
+            if nodes_num < max_nodes:
+                edges = np.pad(edges, ((0, max_nodes - nodes_num), (0, max_nodes - nodes_num)), 'constant',constant_values=0)
+                while len(nodes) < max_nodes:
+                    nodes.append('none')
+            print(f'edges size: {edges.shape}, nodes size: {len(nodes)}')
             return  edges,nodes
         
         def get_nasbench_201_val(idx):
@@ -766,10 +778,12 @@ class Dataset(InMemoryDataset):
 
         with tqdm(total = len_data) as pbar:
             active_nodes = set()
-            file_path = '/home/stud/hanzhang/nasbenchDiT/graph_dit/nasbench-201-graph.json'
+            file_path = '/nfs/data3/hanzhang/nasbenchDiT/graph_dit/nasbench-201-graph.json'
             with open(file_path, 'r') as f:
                 graph_list = json.load(f)
             i = 0
+            flex_graph_list = []
+            flex_graph_path = '/nfs/data3/hanzhang/nasbenchDiT/graph_dit/flex-nasbench201-graph.json'
             for graph in graph_list:
                 # arch_info = self.api.query_meta_info_by_index(i)
                 # results = self.api.query_by_index(i, 'cifar100')
@@ -784,6 +798,16 @@ class Dataset(InMemoryDataset):
                         active_nodes.add(op)
                 
                 data = graph_to_graph_data((adj_matrix, ops)) 
+                # with open(flex_graph_path, 'a') as f:
+                #     flex_graph = {
+                #         'adj_matrix': adj_matrix,
+                #         'ops': ops,
+                #     }
+                #     json.dump(flex_graph, f)
+                flex_graph_list.append({
+                    'adj_matrix':adj_matrix,
+                    'ops': ops,
+                })
                 if i < 3:
                     print(f"i={i}, data={data}")
                     with open(f'{i}.json', 'w') as f:
@@ -792,7 +816,17 @@ class Dataset(InMemoryDataset):
                         f.write(str(data.edge_attr))
                 data_list.append(data)
 
-                new_adj, new_ops = generate_flex_adj_mat(ori_nodes=ori_nodes, ori_edges=ori_adj, max_nodes=12, min_nodes=8,  random_ratio=0.5)
+                new_adj, new_ops = generate_flex_adj_mat(ori_nodes=ori_nodes, ori_edges=ori_adj, max_nodes=12, min_nodes=9,  random_ratio=0.5)
+                flex_graph_list.append({
+                    'adj_matrix':new_adj.tolist(),
+                    'ops': new_ops,
+                })
+                # with open(flex_graph_path, 'w') as f:
+                #     flex_graph = {
+                #         'adj_matrix': new_adj.tolist(),
+                #         'ops': new_ops,
+                #     }
+                #     json.dump(flex_graph, f)
                 data_list.append(graph_to_graph_data((new_adj, new_ops)))
                
                 # graph_list.append({
@@ -838,6 +872,8 @@ class Dataset(InMemoryDataset):
                 graph['ops'] = ops
         with open(f'nasbench-201-graph.json', 'w') as f:
             json.dump(graph_list, f)
+        with open(flex_graph_path, 'w') as f:
+            json.dump(flex_graph_list, f)
             
         torch.save(self.collate(data_list), self.processed_paths[0])
 
@@ -1034,8 +1070,8 @@ def parse_architecture_string(arch_str, padding=0):
             assert idx == steps_coding[cont]
             cont += 1
             nodes.append(n)
-    ori_nodes = nodes.copy()
     nodes.append('output')  # Add output node
+    ori_nodes = nodes.copy()
     if padding > 0:
         for i in range(padding):
             nodes.append('none')
@@ -1048,7 +1084,7 @@ def parse_architecture_string(arch_str, padding=0):
     # print(nodes)
     # print(adj_mat)
     # print(len(adj_mat))
-
+    # print(f'len(ori_nodes): {len(ori_nodes)}, len(nodes): {len(nodes)}')
     return nodes, adj_mat, ori_nodes, ori_adj_mat
 
 def create_adj_matrix_and_ops(nodes, edges):
@@ -1091,6 +1127,7 @@ class DataInfos(AbstractDatasetInfos):
 
             adj_ops_pairs = []
             for item in data:
+                print(item)
                 adj_matrix = np.array(item['adj_matrix'])
                 ops = item['ops']
                 ops = [op_type[op] for op in ops]
@@ -1111,12 +1148,12 @@ class DataInfos(AbstractDatasetInfos):
             #         ops_type[op] = len(ops_type)
             # len_ops.add(len(ops))
             # graphs.append((adj_matrix, ops))
-        graphs = read_adj_ops_from_json(f'/home/stud/hanzhang/nasbenchDiT/graph_dit/nasbench-201-graph.json')
+        graphs = read_adj_ops_from_json(f'/nfs/data3/hanzhang/nasbenchDiT/graph_dit/flex-nasbench201-graph.json')
 
         # check first five graphs
         for i in range(5):
             print(f'graph {i} : {graphs[i]}')
-        print(f'ops_type: {ops_type}')
+        # print(f'ops_type: {ops_type}')
 
         meta_dict = new_graphs_to_json(graphs, 'nasbench-201')
         self.base_path = base_path
@@ -1325,11 +1362,11 @@ def compute_meta(root, source_name, train_index, test_index):
         'transition_E': tansition_E.tolist(),
         }
 
-    with open(f'/home/stud/hanzhang/nasbenchDiT/graph_dit/nasbench201.meta.json', "w") as f:
+    with open(f'/nfs/data3/hanzhang/nasbenchDiT/graph_dit/nasbench201.meta.json', "w") as f:
         json.dump(meta_dict, f)
     
     return meta_dict
 
 
 if __name__ == "__main__":
-    dataset = Dataset(source='nasbench', root='/home/stud/hanzhang/nasbenchDiT/graph-dit', target_prop='Class', transform=None)
+    dataset = Dataset(source='nasbench', root='/nfs/data3/hanzhang/nasbenchDiT/graph-dit', target_prop='Class', transform=None)