使用Pytorch Geometric進行鏈接預(yù)測的實現(xiàn)代碼
PyTorch Geometric (PyG)是構(gòu)建圖神經(jīng)網(wǎng)絡(luò)模型和實驗各種圖卷積的主要工具。在本文中我們將通過鏈接預(yù)測來對其進行介紹。
鏈接預(yù)測答了一個問題:哪兩個節(jié)點應(yīng)該相互鏈接?我們將通過執(zhí)行“轉(zhuǎn)換分割”，為建模準(zhǔn)備數(shù)據(jù)。為批處理準(zhǔn)備專用的圖數(shù)據(jù)加載器。在Torch Geometric中構(gòu)建一個模型，使用PyTorch Lightning進行訓(xùn)練，并檢查模型的性能。
庫準(zhǔn)備
Torch 這個就不用多介紹了
Torch Geometric 圖形神經(jīng)網(wǎng)絡(luò)的主要庫，也是本文介紹的重點
PyTorch Lightning 用于訓(xùn)練、調(diào)優(yōu)和驗證模型。它簡化了訓(xùn)練的操作
Sklearn Metrics和Torchmetrics 用于檢查模型的性能。
PyTorch Geometric有一些特定的依賴關(guān)系，如果你安裝有問題，請參閱其官方文檔。
數(shù)據(jù)準(zhǔn)備
我們將使用Cora ML引文數(shù)據(jù)集。數(shù)據(jù)集可以通過Torch Geometric訪問。
data = tg.datasets.CitationFull(root="data", name="Cora_ML")
默認(rèn)情況下，Torch Geometric數(shù)據(jù)集可以返回多個圖形。我們看看單個圖是什么樣子的
data[0] > Data(x=[2995, 2879], edge_index=[2, 16316], y=[2995])
這里的 X是節(jié)點的特征。edge_index是2 x (n條邊)矩陣(第一維= 2，被解釋為:第0行-源節(jié)點/“發(fā)送方”，第1行-目標(biāo)節(jié)點/“接收方”)。
鏈接拆分
我們將從拆分?jǐn)?shù)據(jù)集中的鏈接開始。使用20%的圖鏈接作為驗證集，10%作為測試集。這里不會向訓(xùn)練數(shù)據(jù)集中添加負(fù)樣本，因為這樣的負(fù)鏈接將由批處理數(shù)據(jù)加載器實時創(chuàng)建。
一般來說，負(fù)采樣會創(chuàng)建“假”樣本(在我們的例子中是節(jié)點之間的鏈接)，因此模型學(xué)習(xí)如何區(qū)分真實和虛假的鏈接。負(fù)抽樣基于抽樣的理論和數(shù)學(xué)，具有一些很好的統(tǒng)計性質(zhì)。
首先:讓我們創(chuàng)建一個鏈接拆分對象。
link_splitter = tg.transforms.RandomLinkSplit(
num_val=0.2,
num_test=0.1,
add_negative_train_samples=False,
disjoint_train_ratio=0.8)
disjoint_train_ratio調(diào)節(jié)在“監(jiān)督”階段將使用多少條邊作為訓(xùn)練信息。剩余的邊將用于消息傳遞(網(wǎng)絡(luò)中的信息傳輸階段)。
圖神經(jīng)網(wǎng)絡(luò)中至少有兩種分割邊的方法:歸納分割和傳導(dǎo)分割。轉(zhuǎn)換方法假設(shè)GNN需要從圖結(jié)構(gòu)中學(xué)習(xí)結(jié)構(gòu)模式。在歸納設(shè)置中，可以使用節(jié)點/邊緣標(biāo)簽進行學(xué)習(xí)。本文最后有兩篇論文詳細(xì)討論了這些概念，并進行了額外的形式化:([1]，[3])。
train_g, val_g, test_g = link_splitter(data[0]) > Data(x=[2995, 2879], edge_index=[2, 2285], y=[2995], edge_label=[9137], edge_label_index=[2, 9137])
在這個操作之后，我們有了一些新的屬性:
edge_label ：描述邊緣是否為真/假。這是我們想要預(yù)測的。
edge_label_index 是一個2 x NUM EDGES矩陣，用于存儲節(jié)點鏈接。
讓我們看看樣本的分布
th.unique(train_g.edge_label, return_counts=True) > (tensor([1.]), tensor([9137])) th.unique(val_g.edge_label, return_counts=True) > (tensor([0., 1.]), tensor([3263, 3263])) th.unique(val_g.edge_label, return_counts=True) > (tensor([0., 1.]), tensor([3263, 3263]))
對于訓(xùn)練數(shù)據(jù)沒有負(fù)邊(我們將訓(xùn)練時創(chuàng)建它們)，對于val/測試集——已經(jīng)以50:50的比例有了一些“假”鏈接。
模型
現(xiàn)在我們可以在使用GNN進行模型的構(gòu)建了一個
class GNN(nn.Module):
def __init__(
self,
dim_in: int,
conv_sizes: Tuple[int, ...],
act_f: nn.Module = th.relu,
dropout: float = 0.1,
*args,
**kwargs):
super().__init__()
self.dim_in = dim_in
self.dim_out = conv_sizes[-1]
self.dropout = dropout
self.act_f = act_f
last_in = dim_in
layers = []
# Here we build subsequent graph convolutions.
for conv_sz in conv_sizes:
# Single graph convolution layer
conv = tgnn.SAGEConv(in_channels=last_in, out_channels=conv_sz, *args, **kwargs)
last_in = conv_sz
layers.append(conv)
self.layers = nn.ModuleList(layers)
def forward(self, x: th.Tensor, edge_index: th.Tensor) -> th.Tensor:
h = x
# For every graph convolution in the network...
for conv in self.layers:
# ... perform node embedding via message passing
h = conv(h, edge_index)
h = self.act_f(h)
if self.dropout:
h = nn.functional.dropout(h, p=self.dropout, training=self.training)
return h
這個模型中值得注意的部分是一組圖卷積——在我們的例子中是SAGEConv。SAGE卷積的正式定義為:
v是當(dāng)前節(jié)點，節(jié)點v的N(v)個鄰居。要了解更多關(guān)于這種卷積類型的信息，請查看GraphSAGE[1]的原始論文
讓我們檢查一下模型是否可以使用準(zhǔn)備好的數(shù)據(jù)進行預(yù)測。這里PyG模型的輸入是節(jié)點特征X的矩陣和定義edge_index的鏈接。
gnn = GNN(train_g.x.size()[1], conv_sizes=[512, 256, 128])
with th.no_grad():
out = gnn(train_g.x, train_g.edge_index)
out
> tensor([[0.0000, 0.0000, 0.0051, ..., 0.0997, 0.0000, 0.0000],
[0.0107, 0.0000, 0.0576, ..., 0.0651, 0.0000, 0.0000],
[0.0000, 0.0000, 0.0102, ..., 0.0973, 0.0000, 0.0000],
...,
[0.0000, 0.0000, 0.0549, ..., 0.0671, 0.0000, 0.0000],
[0.0000, 0.0000, 0.0166, ..., 0.0000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.0034, ..., 0.1111, 0.0000, 0.0000]])
我們模型的輸出是一個維度為:N個節(jié)點x嵌入大小的節(jié)點嵌入矩陣。
PyTorch Lightning
PyTorch Lightning主要用作訓(xùn)練，但是這里我們在GNN的輸出后面增加了一個Linear層做為預(yù)測是否鏈接的輸出頭。
class LinkPredModel(pl.LightningModule):
def __init__(
self,
dim_in: int,
conv_sizes: Tuple[int, ...],
act_f: nn.Module = th.relu,
dropout: float = 0.1,
lr: float = 0.01,
*args, **kwargs):
super().__init__()
# Our inner GNN model
self.gnn = GNN(dim_in, conv_sizes=conv_sizes, act_f=act_f, dropout=dropout)
# Final prediction model on links.
self.lin_pred = nn.Linear(self.gnn.dim_out, 1)
self.lr = lr
def forward(self, x: th.Tensor, edge_index: th.Tensor) -> th.Tensor:
# Step 1: make node embeddings using GNN.
h = self.gnn(x, edge_index)
# Take source nodes embeddings- senders
h_src = h[edge_index[0, :]]
# Take target node embeddings - receivers
h_dst = h[edge_index[1, :]]
# Calculate the product between them
src_dst_mult = h_src * h_dst
# Apply non-linearity
out = self.lin_pred(src_dst_mult)
return out
def _step(self, batch: th.Tensor, phase: str='train') -> th.Tensor:
yhat_edge = self(batch.x, batch.edge_label_index).squeeze()
y = batch.edge_label
loss = nn.functional.binary_cross_entropy_with_logits(input=yhat_edge, target=y)
f1 = tm.functional.f1_score(preds=yhat_edge, target=y, task='binary')
prec = tm.functional.precision(preds=yhat_edge, target=y, task='binary')
recall = tm.functional.recall(preds=yhat_edge, target=y, task='binary')
# Watch for logging here - we need to provide batch_size, as (at the time of this implementation)
# PL cannot understand the batch size.
self.log(f"{phase}_f1", f1, batch_size=batch.edge_label_index.shape[1])
self.log(f"{phase}_loss", loss, batch_size=batch.edge_label_index.shape[1])
self.log(f"{phase}_precision", prec, batch_size=batch.edge_label_index.shape[1])
self.log(f"{phase}_recall", recall, batch_size=batch.edge_label_index.shape[1])
return loss
def training_step(self, batch, batch_idx):
return self._step(batch)
def validation_step(self, batch, batch_idx):
return self._step(batch, "val")
def test_step(self, batch, batch_idx):
return self._step(batch, "test")
def predict_step(self, batch):
x, edge_index = batch
return self(x, edge_index)
def configure_optimizers(self):
return th.optim.Adam(self.parameters(), lr=self.lr)
PyTorch Lightning的作用是幫我們簡化了訓(xùn)練的步驟，我們只需要配置一些函數(shù)即可，我們可以使用以下命令測試模型是否可用
model = LinkPredModel(val_g.x.size()[1], conv_sizes=[512, 256, 128])
with th.no_grad():
out = model.predict_step((val_g.x, val_g.edge_label_index))
訓(xùn)練
對于訓(xùn)練的步驟，需要特殊處理的是數(shù)據(jù)加載器。
圖數(shù)據(jù)需要特殊處理——尤其是鏈接預(yù)測。PyG有一些專門的數(shù)據(jù)加載器類，它們負(fù)責(zé)正確地生成批處理。我們將使用:tg.loader.LinkNeighborLoader，它接受以下輸入:
要批量加載的數(shù)據(jù)(圖)。num_neighbors 每個節(jié)點在一次“跳”期間加載的最大鄰居數(shù)量。指定鄰居數(shù)目的列表1 - 2 - 3 -…-K。對于非常大的圖形特別有用。
edge_label_index 哪個屬性已經(jīng)指示了真/假鏈接。
neg_sampling_ratio -負(fù)樣本與真實樣本的比例。
train_loader = tg.loader.LinkNeighborLoader(
train_g,
num_neighbors=[-1, 10, 5],
batch_size=128,
edge_label_index=train_g.edge_label_index,
# "on the fly" negative sampling creation for batch
neg_sampling_ratio=0.5
)
val_loader = tg.loader.LinkNeighborLoader(
val_g,
num_neighbors=[-1, 10, 5],
batch_size=128,
edge_label_index=val_g.edge_label_index,
edge_label=val_g.edge_label,
# negative samples for val set are done already as ground-truth
neg_sampling_ratio=0.0
)
test_loader = tg.loader.LinkNeighborLoader(
test_g,
num_neighbors=[-1, 10, 5],
batch_size=128,
edge_label_index=test_g.edge_label_index,
edge_label=test_g.edge_label,
# negative samples for test set are done already as ground-truth
neg_sampling_ratio=0.0
)
下面就是訓(xùn)練模型
model = LinkPredModel(val_g.x.size()[1], conv_sizes=[512, 256, 128]) trainer = pl.Trainer(max_epochs=20, log_every_n_steps=5) # Validate before training - we will see results of untrained model. trainer.validate(model, val_loader) # Train the model trainer.fit(model=model, train_dataloaders=train_loader, val_dataloaders=val_loader)
試驗數(shù)據(jù)核對，查看分類報告和ROC曲線。
with th.no_grad():
yhat_test_proba = th.sigmoid(model(test_g.x, test_g.edge_label_index)).squeeze()
yhat_test_cls = yhat_test_proba >= 0.5
print(classification_report(y_true=test_g.edge_label, y_pred=yhat_test_cls))
結(jié)果看起來還不錯:
precision recall f1-score support
0.0 0.68 0.70 0.69 1631
1.0 0.69 0.66 0.68 1631
accuracy 0.68 3262
macro avg 0.68 0.68 0.68 3262
weighted avg 0.68 0.68 0.68 3262
ROC曲線也不錯
我們訓(xùn)練的模型并不特別復(fù)雜，也沒有經(jīng)過精心調(diào)整，但它完成了工作。當(dāng)然這只是一個為了演示使用的小型數(shù)據(jù)集。
總結(jié)
圖神經(jīng)網(wǎng)絡(luò)盡管看起來很復(fù)雜，但是PyTorch Geometric為我們提供了一個很好的解決方案。我們可以直接使用其中內(nèi)置的模型實現(xiàn)，這方便了我們使用和簡化了入門的門檻。
本文代碼：
https://avoid.overfit.cn/post/e14c4369776243d68c22c4a2a0346db2
以上就是使用Pytorch Geometric進行鏈接預(yù)測的實現(xiàn)代碼的詳細(xì)內(nèi)容，更多關(guān)于Pytorch Geometric鏈接預(yù)測的資料請關(guān)注腳本之家其它相關(guān)文章！
相關(guān)文章
你是否常常感到桌面文件太亂,各類文檔、圖片、視頻混雜一團,找個文件如大海撈針,桌面圖標(biāo)密密麻麻,各種新建文件夾(1)、圖片001、合同終版(最終).pdf堆積如山,今天,我就將手把手教你如何利用Python實現(xiàn)全自動文件整理,需要的朋友可以參考下2025-11-11
在本篇文章里小編給大家整理的是關(guān)于python解包用法詳解內(nèi)容，有需要的朋友們可以跟著一起學(xué)習(xí)下。2021-02-02
這篇文章主要介紹了在主流系統(tǒng)之上安裝Pygame的方法，本文通過實例圖文相結(jié)合給大家介紹的非常詳細(xì)，對大家的學(xué)習(xí)或工作具有一定的參考借鑒價值,需要的朋友可以參考下2020-05-05
從Node.js+TypeScript無縫切換到Python的完整過程實錄
隨著技術(shù)的發(fā)展,越來越多的項目需要同時使用多種編程語言進行開發(fā),Python和Node.js都是非常流行的編程語言,分別有著各自的優(yōu)勢和適用場景,這篇文章主要介紹了從Node.js+TypeScript無縫切換到Python的完整過程,需要的朋友可以參考下2026-06-06
在自然語言處理領(lǐng)域,文本分詞是最基礎(chǔ)也是最關(guān)鍵的技術(shù)環(huán)節(jié),本文將深入解析Python文本分詞技術(shù)體系,希望對大家有一定的幫助2025-08-08
這篇文章主要介紹了Python運維開發(fā)之psutil庫的使用，psutil能夠輕松實現(xiàn)獲取系統(tǒng)運行的進程和系統(tǒng)利用率。小編覺得挺不錯的，現(xiàn)在分享給大家，也給大家做個參考。一起跟隨小編過來看看吧2018-10-10
本文詳細(xì)且清晰地講解了Python中錯誤和異常的概念及其處理方式，通過具體案例展示try...except、try...finally、with...等句式的具體用法，期望能幫助到對此感到迷惑的初學(xué)者2021-10-10
Django提供了兩個非常有用的工具：F對象和Q對象，方便了在一些特殊場景下的查詢過程，這篇文章主要介紹了Django?ORM?F對象和Q對象查詢,需要的朋友可以參考下2022-10-10
python?tkinter實現(xiàn)學(xué)生信息管理系統(tǒng)
這篇文章主要為大家詳細(xì)介紹了python?tkinter實現(xiàn)學(xué)生信息管理系統(tǒng)，文中示例代碼介紹的非常詳細(xì)，具有一定的參考價值，感興趣的小伙伴們可以參考一下2022-02-02
這篇文章主要介紹了python UNIX_TIMESTAMP時間處理方法,結(jié)合實例形式分析了Python針對UNIX_TIMESTAMP時間的常見運算技巧,需要的朋友可以參考下2016-04-04