ML model/Graph Neural Network Model
(GCN) Graph Convolution Network
코딩소비
2024. 1. 25. 13:39
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)