(GCN) Graph Convolution Network

2024. 1. 25. 13:39ML model/Graph Neural Network Model

input : edges, labels

from torch_geometric.nn.models import LightGCN

edges = [
	[node_id1, node_id2],
    [node_id1, node_id3],
    [node_id2, node_id10],
    ...
	]
    
labels = [1,0,0,1,1,1, ...] # for link prediction

model = LightGCN(num_nodes=n_node, **kwargs)

pred = model(edges)
loss = model.link_pred_loss(pred=pred, edge_label=labels)

prob = model.predict_link(edge_index=edges, prob=True)
acc = accuracy_score(y_true=labels, y_pred=prob > 0.5)
auc = roc_auc_score(y_true=labels, y_score=prob)

 

 

 

데이터 전처리

#id2index는 data.userID와 data.assessmentItemID 각각 0부터 오름차순 id를 mapping한 dictionary

edge, label = [], []
for user, item, acode in zip(data.userID, data.assessmentItemID, data.answerCode):
    uid, iid = id2index[user], id2index[item]
    edge.append([uid, iid])
    label.append(acode)

edge = torch.LongTensor(edge).T
label = torch.LongTensor(label)