-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathGCN_without_edgeFeatures.py
205 lines (164 loc) · 8 KB
/
GCN_without_edgeFeatures.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
import os
import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import networkx as nx
import stellargraph as sg
from stellargraph.mapper import PaddedGraphGenerator
from stellargraph.layer import DeepGraphCNN
from stellargraph import StellarGraph
import pickle
import warnings
# Filter out the UserWarning from StellarGraph
warnings.filterwarnings("ignore", category=UserWarning, module="stellargraph")
from sklearn.metrics import confusion_matrix, classification_report
from tensorflow.keras import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.layers import Dense, Conv1D, MaxPool1D, Dropout, Flatten
from tensorflow.keras.losses import sparse_categorical_crossentropy # Use sparse categorical cross-entropy
import tensorflow as tf
from sklearn import model_selection
from tensorflow.keras.models import load_model
from stellargraph.data import EdgeSplitter
from stellargraph.mapper import GraphSAGENodeGenerator
from stellargraph.layer import GraphSAGE
from stellargraph.mapper import PaddedGraphGenerator
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from tensorflow.keras import layers, optimizers, losses, metrics, models
from tensorflow.keras.callbacks import EarlyStopping
# Load data from JSON files
data_dir = 'C:\\Users\\Saqib\\PycharmProjects\\GAGE\\output_AED'
data = []
labels = []
max_nodes = 128 # Maximum number of nodes for each graph
for class_folder in os.listdir(data_dir):
print(class_folder)
if os.path.isdir(os.path.join(data_dir, class_folder)):
class_label = class_folder
class_folder_path = os.path.join(data_dir, class_folder)
for json_file in os.listdir(class_folder_path):
if json_file.endswith('.json'):
with open(os.path.join(class_folder_path, json_file), 'r') as file:
graph_data = json.load(file)
# Check if JSON data has nodes and edges
if "Nodes" not in graph_data or "Edges" not in graph_data:
print(f"JSON file {json_file} has no nodes or edges, skipping...")
continue
# Extract only CPID values from nodes
nodes = [node["CPID"] for node in graph_data["Nodes"]]
if not nodes:
print(f"JSON file {json_file} has no node data, skipping...")
continue
# Limit the number of nodes to max_nodes
nodes = nodes[:max_nodes]
# Filter edges connected to nodes within the first max_nodes
edges = [edge for edge in graph_data["Edges"] if
edge["Source"] in nodes or edge["Destination"] in nodes]
# Extract edge features
edge_features = []
for edge in edges:
edge_features.append([
edge.get("is_consequent", False),
edge.get("is_conditional", False),
edge.get("intra_edge", False),
edge.get("external_edge", False)
])
# Extract node features based on CPID values
node_features = {node["CPID"]: node["AED_features"] for node in graph_data["Nodes"] if
node["CPID"] in nodes}
graph_data["Nodes"] = nodes
graph_data["Edges"] = edges
graph_data["NodeFeatures"] = node_features
graph_data["EdgeFeatures"] = edge_features
data.append(graph_data)
labels.append(class_label)
# Encode class labels
num_classes = len(set(labels)) # Calculate the number of unique classes
class_to_index = {c: i for i, c in enumerate(set(labels))}
labels = [class_to_index[label] for label in labels]
# Create a StellarGraph object for each graph in your dataset
# Define a list to store individual graph data
graphs = []
# Iterate through your data and create a StellarGraph for each graph
for graph_data in data:
# Create an empty graph
G = nx.Graph()
# Add nodes to the graph
for node_id in graph_data["Nodes"]:
G.add_node(node_id, features=np.array(graph_data["NodeFeatures"][node_id][0]))
# Add edges to the graph
i = 0
for edge in graph_data["Edges"]:
source, target = edge["Source"], edge["Destination"]
# fe = graph_data["EdgeFeatures"][i]
# int_fe = int_list = [int(x) for x in fe]
# decimal_fe = sum(b << i for i, b in enumerate(reversed(int_fe)))
# G.add_edge(source, target, Edge_Features=decimal_fe) # Add edge features
G.add_edge(source, target)
i += 1
# Create a StellarGraph from the NetworkX graph
# stellar_graph = sg.StellarGraph.from_networkx(G, node_features="features", edge_type_attr="Edge_Features")
stellar_graph = sg.StellarGraph.from_networkx(G, node_features="features")
graphs.append(stellar_graph)
# Split data into train, validation, and test sets
train_data, test_data, train_labels, test_labels = train_test_split(graphs, labels, test_size=0.2, random_state=42)
train_data, val_data, train_labels, val_labels = train_test_split(train_data, train_labels, test_size=0.2, random_state=42)
# Generator
generator = PaddedGraphGenerator(graphs=train_data)
# Save the generator to a file
with open('models/padded_graph_generator_20231009_GCN2_v1_withoutEdgeFeature.pkl', 'wb') as file:
pickle.dump(generator, file)
k = 35
layer_sizes = [64, 32, 16, num_classes] # Adjust the last layer size for multi-class
dgcnn_model = DeepGraphCNN(
layer_sizes=layer_sizes,
activations=["tanh", "tanh", "tanh", "softmax"], # Use softmax for multi-class
k=k,
bias=False,
generator=generator,
)
x_inp, x_out = dgcnn_model.in_out_tensors()
x_out = Conv1D(filters=16, kernel_size=sum(layer_sizes), strides=sum(layer_sizes))(x_out)
x_out = MaxPool1D(pool_size=2)(x_out)
x_out = Conv1D(filters=32, kernel_size=5, strides=1)(x_out)
x_out = Flatten()(x_out)
x_out = Dense(units=128, activation="relu")(x_out)
x_out = Dropout(rate=0.5)(x_out)
predictions = Dense(units=num_classes, activation="softmax")(x_out) # Use softmax for multi-class
model = Model(inputs=x_inp, outputs=predictions)
model.compile(
optimizer=Adam(lr=0.0001), loss=sparse_categorical_crossentropy, metrics=["acc"],
)
train_gen = generator.flow(train_data, train_labels, batch_size=50, symmetric_normalization=False)
val_gen = generator.flow(val_data, val_labels, batch_size=1, symmetric_normalization=False)
test_gen = generator.flow(test_data, test_labels, batch_size=1, symmetric_normalization=False)
early_stopping = EarlyStopping(monitor='val_loss', patience=50, restore_best_weights=True)
history = model.fit(
train_gen, epochs=200, verbose=1, validation_data=val_gen, shuffle=True, callbacks=[early_stopping]
)
# Save the trained model
# Define a function to save the model
def save_model(model, model_name):
model.save(model_name)
model_name = "models/GCN_20231009_GCN2_v1_withoutEdgeFeature.h5"
save_model(model, model_name)
sg.utils.plot_history(history)
plt.show()
# Evaluate the model on the test set
test_metrics = model.evaluate(test_gen)
print("\nTest Set Metrics:")
for name, val in zip(model.metrics_names, test_metrics):
print("\t{}: {:0.4f}".format(name, val))
# Make predictions on the test set
Y_pred = model.predict(test_gen)
Y_pred_classes = np.argmax(Y_pred, axis=1)
# Classification report
class_names = [index for index, _ in sorted(class_to_index.items(), key=lambda x: x[1])]
report = classification_report(test_labels, Y_pred_classes, target_names=class_names)
print(report)
# Confusion matrix
confusion = confusion_matrix(test_labels, Y_pred_classes)
print(confusion)
print(class_to_index)