使用Pytorch Geometric進(jìn)行鏈接預(yù)測(cè)的實(shí)現(xiàn)代碼

更新時(shí)間：2023年10月23日 08:51:49 作者：deephub

PyTorch Geometric (PyG)是構(gòu)建圖神經(jīng)網(wǎng)絡(luò)模型和實(shí)驗(yàn)各種圖卷積的主要工具,在本文中我們將通過(guò)鏈接預(yù)測(cè)來(lái)對(duì)其進(jìn)行介紹,文中有詳細(xì)的代碼示例供大家參考,需要的朋友可以參考下

鏈接預(yù)測(cè)答了一個(gè)問(wèn)題:哪兩個(gè)節(jié)點(diǎn)應(yīng)該相互鏈接?我們將通過(guò)執(zhí)行“轉(zhuǎn)換分割”，為建模準(zhǔn)備數(shù)據(jù)。為批處理準(zhǔn)備專用的圖數(shù)據(jù)加載器。在Torch Geometric中構(gòu)建一個(gè)模型，使用PyTorch Lightning進(jìn)行訓(xùn)練，并檢查模型的性能。

庫(kù)準(zhǔn)備

Torch 這個(gè)就不用多介紹了

Torch Geometric 圖形神經(jīng)網(wǎng)絡(luò)的主要庫(kù)，也是本文介紹的重點(diǎn)

PyTorch Lightning 用于訓(xùn)練、調(diào)優(yōu)和驗(yàn)證模型。它簡(jiǎn)化了訓(xùn)練的操作

Sklearn Metrics和Torchmetrics 用于檢查模型的性能。

PyTorch Geometric有一些特定的依賴關(guān)系，如果你安裝有問(wèn)題，請(qǐng)參閱其官方文檔。

數(shù)據(jù)準(zhǔn)備

我們將使用Cora ML引文數(shù)據(jù)集。數(shù)據(jù)集可以通過(guò)Torch Geometric訪問(wèn)。

 data = tg.datasets.CitationFull(root="data", name="Cora_ML")

默認(rèn)情況下，Torch Geometric數(shù)據(jù)集可以返回多個(gè)圖形。我們看看單個(gè)圖是什么樣子的

 data[0]
 > Data(x=[2995, 2879], edge_index=[2, 16316], y=[2995])

這里的 X是節(jié)點(diǎn)的特征。edge_index是2 x (n條邊)矩陣(第一維= 2，被解釋為:第0行-源節(jié)點(diǎn)/“發(fā)送方”，第1行-目標(biāo)節(jié)點(diǎn)/“接收方”)。

鏈接拆分

我們將從拆分?jǐn)?shù)據(jù)集中的鏈接開(kāi)始。使用20%的圖鏈接作為驗(yàn)證集，10%作為測(cè)試集。這里不會(huì)向訓(xùn)練數(shù)據(jù)集中添加負(fù)樣本，因?yàn)檫@樣的負(fù)鏈接將由批處理數(shù)據(jù)加載器實(shí)時(shí)創(chuàng)建。

一般來(lái)說(shuō)，負(fù)采樣會(huì)創(chuàng)建“假”樣本(在我們的例子中是節(jié)點(diǎn)之間的鏈接)，因此模型學(xué)習(xí)如何區(qū)分真實(shí)和虛假的鏈接。負(fù)抽樣基于抽樣的理論和數(shù)學(xué)，具有一些很好的統(tǒng)計(jì)性質(zhì)。

首先:讓我們創(chuàng)建一個(gè)鏈接拆分對(duì)象。

 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é)點(diǎn)/邊緣標(biāo)簽進(jìn)行學(xué)習(xí)。本文最后有兩篇論文詳細(xì)討論了這些概念，并進(jìn)行了額外的形式化:([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])

在這個(gè)操作之后，我們有了一些新的屬性:

edge_label ：描述邊緣是否為真/假。這是我們想要預(yù)測(cè)的。

edge_label_index 是一個(gè)2 x NUM EDGES矩陣，用于存儲(chǔ)節(jié)點(diǎn)鏈接。

讓我們看看樣本的分布

 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]))

對(duì)于訓(xùn)練數(shù)據(jù)沒(méi)有負(fù)邊(我們將訓(xùn)練時(shí)創(chuàng)建它們)，對(duì)于val/測(cè)試集——已經(jīng)以50:50的比例有了一些“假”鏈接。

模型

現(xiàn)在我們可以在使用GNN進(jìn)行模型的構(gòu)建了一個(gè)

 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

這個(gè)模型中值得注意的部分是一組圖卷積——在我們的例子中是SAGEConv。SAGE卷積的正式定義為:

v是當(dāng)前節(jié)點(diǎn)，節(jié)點(diǎn)v的N(v)個(gè)鄰居。要了解更多關(guān)于這種卷積類型的信息，請(qǐng)查看GraphSAGE[1]的原始論文

讓我們檢查一下模型是否可以使用準(zhǔn)備好的數(shù)據(jù)進(jìn)行預(yù)測(cè)。這里PyG模型的輸入是節(jié)點(diǎn)特征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]])

我們模型的輸出是一個(gè)維度為:N個(gè)節(jié)點(diǎn)x嵌入大小的節(jié)點(diǎn)嵌入矩陣。

PyTorch Lightning

PyTorch Lightning主要用作訓(xùn)練，但是這里我們?cè)贕NN的輸出后面增加了一個(gè)Linear層做為預(yù)測(cè)是否鏈接的輸出頭。

 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的作用是幫我們簡(jiǎn)化了訓(xùn)練的步驟，我們只需要配置一些函數(shù)即可，我們可以使用以下命令測(cè)試模型是否可用

 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)練

對(duì)于訓(xùn)練的步驟，需要特殊處理的是數(shù)據(jù)加載器。

圖數(shù)據(jù)需要特殊處理——尤其是鏈接預(yù)測(cè)。PyG有一些專門的數(shù)據(jù)加載器類，它們負(fù)責(zé)正確地生成批處理。我們將使用:tg.loader.LinkNeighborLoader，它接受以下輸入:

要批量加載的數(shù)據(jù)(圖)。num_neighbors 每個(gè)節(jié)點(diǎn)在一次“跳”期間加載的最大鄰居數(shù)量。指定鄰居數(shù)目的列表1 - 2 - 3 -…-K。對(duì)于非常大的圖形特別有用。

edge_label_index 哪個(gè)屬性已經(jīng)指示了真/假鏈接。

neg_sampling_ratio -負(fù)樣本與真實(shí)樣本的比例。

 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)

試驗(yàn)數(shù)據(jù)核對(duì)，查看分類報(bào)告和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é)果看起來(lái)還不錯(cuò):

               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曲線也不錯(cuò)