forked from Eshe/competition-vd
72 lines
2.2 KiB
Python
Executable File
72 lines
2.2 KiB
Python
Executable File
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
import networkx as nx
|
|
from model_magic import GNNClassifier
|
|
from utils_and_helpers import load_data, compute_metrics
|
|
|
|
def create_graph(function):
|
|
G = nx.Graph()
|
|
lines = function.split('\n')
|
|
for i, line in enumerate(lines):
|
|
G.add_node(i, features=torch.randn(64))
|
|
if i > 0:
|
|
G.add_edge(i-1, i)
|
|
return G
|
|
|
|
def graph_to_tensor(G):
|
|
num_nodes = G.number_of_nodes()
|
|
adj = torch.zeros((num_nodes, num_nodes))
|
|
features = torch.stack([G.nodes[i]['features'] for i in range(num_nodes)])
|
|
|
|
for edge in G.edges():
|
|
adj[edge[0]][edge[1]] = 1
|
|
adj[edge[1]][edge[0]] = 1
|
|
|
|
return features, adj
|
|
|
|
def train_gnn(model, graphs, labels, num_epochs):
|
|
criterion = nn.CrossEntropyLoss()
|
|
optimizer = optim.Adam(model.parameters(), lr=0.01)
|
|
|
|
for epoch in range(num_epochs):
|
|
model.train()
|
|
total_loss = 0
|
|
for graph, label in zip(graphs, labels):
|
|
features, adj = graph_to_tensor(graph)
|
|
optimizer.zero_grad()
|
|
output = model(features, adj)
|
|
loss = criterion(output.unsqueeze(0), torch.tensor([label]))
|
|
loss.backward()
|
|
optimizer.step()
|
|
total_loss += loss.item()
|
|
|
|
print(f"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss/len(graphs)}")
|
|
|
|
return model
|
|
|
|
def main():
|
|
functions, labels = load_data('output_dump/processed_data')
|
|
graphs = [create_graph(func) for func in functions[:1000]] # Limit to 1000 for demonstration
|
|
|
|
num_classes = len(set(labels))
|
|
model = GNNClassifier(input_dim=64, hidden_dim=32, num_classes=num_classes)
|
|
|
|
trained_model = train_gnn(model, graphs, labels[:1000], num_epochs=10)
|
|
|
|
test_graphs = [create_graph(func) for func in functions[1000:1100]]
|
|
test_labels = labels[1000:1100]
|
|
|
|
model.eval()
|
|
predictions = []
|
|
for graph in test_graphs:
|
|
features, adj = graph_to_tensor(graph)
|
|
output = model(features, adj)
|
|
_, pred = torch.max(output, 0)
|
|
predictions.append(pred.item())
|
|
|
|
metrics = compute_metrics(test_labels, predictions)
|
|
print("GNN Test Metrics:", metrics)
|
|
|
|
if __name__ == "__main__":
|
|
main() |