iicd/Graph-DiT
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update nasbench201 adj_mat and ops mat

Browse Source
This commit is contained in:
mhz 2024-07-04 22:01:36 +02:00
parent d44900c8ba
commit 2cb07838e1
1 changed files with 66 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

96
graph_dit/datasets/dataset.py
View File

@ -41,13 +41,13 @@ op_to_atom = {
}
op_type = {
'input': 0,
'nor_conv_1x1': 1,
'nor_conv_3x3': 2,
'avg_pool_3x3': 3,
'skip_connect': 4,
'output': 5,
'none': 6,
'input': 7
'none': 5,
'output': 6,
}
class DataModule(AbstractDataModule):
def __init__(self, cfg):
@ -130,20 +130,44 @@ class DataModule(AbstractDataModule):
print(self.task, ' dataset len', len(dataset), 'train len', len(train_index), 'val len', len(val_index), 'test len', len(test_index))
return train_index, val_index, test_index, []
def parse_architecture_string(self, arch_str):
stages = arch_str.split('+')
nodes = ['input']
edges = []
# def parse_architecture_string(self, arch_str):
# stages = arch_str.split('+')
# nodes = ['input']
# edges = []
for stage in stages:
operations = stage.strip('|').split('|')
for op in operations:
operation, idx = op.split('~')
idx = int(idx)
edges.append((idx, len(nodes))) # Add edge from idx to the new node
nodes.append(operation)
nodes.append('output') # Add the output node
return nodes, edges
# for stage in stages:
# operations = stage.strip('|').split('|')
# for op in operations:
# operation, idx = op.split('~')
# idx = int(idx)
# edges.append((idx, len(nodes))) # Add edge from idx to the new node
# nodes.append(operation)
# nodes.append('output') # Add the output node
# return nodes, edges
def parse_architecture_string(arch_str):
# print(arch_str)
steps = arch_str.split('+')
nodes = ['input'] # Start with input node
adj_mat = np.array([[0, 1, 1, 0, 1, 0, 0, 0],
[0, 0, 0, 1, 0, 1 ,0 ,0],
[0, 0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 0, 1],
[0, 0, 0, 0, 0, 0, 0, 1],
[0, 0, 0, 0, 0, 0, 0, 1],
[0, 0, 0, 0, 0, 0, 0, 0]])
steps = arch_str.split('+')
steps_coding = ['0', '0', '1', '0', '1', '2']
cont = 0
for step in steps:
step = step.strip('|').split('|')
for node in step:
n, idx = node.split('~')
assert idx == steps_coding[cont]
cont += 1
nodes.append(n)
nodes.append('output') # Add output node
return nodes, adj_mat
# def create_molecule_from_graph(nodes, edges):
def create_molecule_from_graph(self, graph):
@ -182,11 +206,11 @@ class DataModule(AbstractDataModule):
return mol
def arch_str_to_smiles(self, arch_str):
nodes, edges = self.parse_architecture_string(arch_str)
mol = self.create_molecule_from_graph(nodes, edges)
smiles = Chem.MolToSmiles(mol)
return smiles
# def arch_str_to_smiles(self, arch_str):
# nodes, edges = self.parse_architecture_string(arch_str)
# mol = self.create_molecule_from_graph(nodes, edges)
# smiles = Chem.MolToSmiles(mol)
# return smiles
def get_train_graphs(self):
train_graphs = []
@ -684,8 +708,9 @@ class Dataset(InMemoryDataset):
for i in range(len_data):
arch_info = self.api.query_meta_info_by_index(i)
results = self.api.query_by_index(i, 'cifar100')
nodes, edges = parse_architecture_string(arch_info.arch_str)
adj_matrix, ops = create_adj_matrix_and_ops(nodes, edges)
# nodes, edges = parse_architecture_string(arch_info.arch_str)
ops, adj_matrix = parse_architecture_string(arch_info.arch_str)
# adj_matrix, ops = create_adj_matrix_and_ops(nodes, edges)
for op in ops:
if op not in active_nodes:
active_nodes.add(op)
@ -901,15 +926,26 @@ def parse_architecture_string(arch_str):
# print(arch_str)
steps = arch_str.split('+')
nodes = ['input'] # Start with input node
edges = []
for i, step in enumerate(steps):
step = step.strip('|').split('|')
adj_mat = np.array([[0, 1, 1, 0, 1, 0, 0, 0],
[0, 0, 0, 1, 0, 1 ,0 ,0],
[0, 0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 0, 1],
[0, 0, 0, 0, 0, 0, 0, 1],
[0, 0, 0, 0, 0, 0, 0, 1],
[0, 0, 0, 0, 0, 0, 0, 0]])
steps = arch_str.split('+')
steps_coding = ['0', '0', '1', '0', '1', '2']
cont = 0
for step in steps:
step = step.strip('|').split('|')
for node in step:
op, idx = node.split('~')
edges.append((int(idx), i+1)) # i+1 because 0 is input node
nodes.append(op)
n, idx = node.split('~')
assert idx == steps_coding[cont]
cont += 1
nodes.append(n)
nodes.append('output') # Add output node
return nodes, edges
return nodes, adj_mat
def create_adj_matrix_and_ops(nodes, edges):
num_nodes = len(nodes)