???????到目前為止，我們已經(jīng)為節(jié)點(diǎn)分類任務(wù)單獨(dú)以全批方式訓(xùn)練了圖神經(jīng)網(wǎng)絡(luò)。特別是，這意味著每個(gè)節(jié)點(diǎn)的隱藏表示都是并行計(jì)算的，并且可以在下一層中重復(fù)使用。

???????然而，一旦我們想在更大的圖上操作，由于內(nèi)存消耗爆炸，這種方案就不再可行。例如，一個(gè)具有大約1000萬(wàn)個(gè)節(jié)點(diǎn)和128個(gè)隱藏特征維度的圖已經(jīng)為每層消耗了大約5GB的GPU內(nèi)存。

???????因此，最近有一些努力讓圖神經(jīng)網(wǎng)絡(luò)擴(kuò)展到更大的圖。其中一種方法被稱為Cluster-GCN (Chiang et al. (2019)，它基于將圖預(yù)先劃分為子圖，可以在子圖上以小批量的方式進(jìn)行操作。

???????為了展示，讓我們從 Planetoid 節(jié)點(diǎn)分類基準(zhǔn)套件(Yang et al. (2016))中加載PubMed 圖：

import torch
from torch_geometric.datasets import Planetoid
from torch_geometric.transforms import NormalizeFeaturesdataset = Planetoid(root='data/Planetoid', name='PubMed', transform=NormalizeFeatures())print()
print(f'Dataset: {dataset}:')
print('==================')
print(f'Number of graphs: {len(dataset)}')
print(f'Number of features: {dataset.num_features}')
print(f'Number of classes: {dataset.num_classes}')data = dataset[0]  # Get the first graph object.print()
print(data)
print('===============================================================================================================')# Gather some statistics about the graph.
print(f'Number of nodes: {data.num_nodes}')
print(f'Number of edges: {data.num_edges}')
print(f'Average node degree: {data.num_edges / data.num_nodes:.2f}')
print(f'Number of training nodes: {data.train_mask.sum()}')
print(f'Training node label rate: {int(data.train_mask.sum()) / data.num_nodes:.3f}')
print(f'Has isolated nodes: {data.has_isolated_nodes()}')
print(f'Has self-loops: {data.has_self_loops()}')
print(f'Is undirected: {data.is_undirected()}')

???????可以看出，該圖大約有19717個(gè)節(jié)點(diǎn)。雖然這個(gè)數(shù)量的節(jié)點(diǎn)應(yīng)該可以輕松地放入GPU內(nèi)存，但它仍然是一個(gè)很好的例子，可以展示如何在PyTorch Geometric中擴(kuò)展GNN。

???????Cluster-GCN的工作原理是首先基于圖劃分算法將圖劃分為子圖。因此，GNN被限制為僅在其特定子圖內(nèi)進(jìn)行卷積，從而省略了鄰域爆炸的問(wèn)題。

???????然而，在對(duì)圖進(jìn)行分區(qū)后，會(huì)刪除一些鏈接，這可能會(huì)由于有偏差的估計(jì)而限制模型的性能。為了解決這個(gè)問(wèn)題，Cluster-GCN還將類別之間的連接合并到一個(gè)小批量中，這導(dǎo)致了以下隨機(jī)劃分方案：

???????這里，顏色表示每批維護(hù)的鄰接信息（對(duì)于每個(gè)epoch可能不同）。PyTorch Geometric提供了Cluster-GCN算法的兩階段實(shí)現(xiàn)：

1. ClusterData 將一個(gè) Data 對(duì)象轉(zhuǎn)換為包含num_parts分區(qū)的子圖的數(shù)據(jù)集。
2. 給定一個(gè)用戶定義的batch_size, ClusterLoader 實(shí)現(xiàn)隨機(jī)劃分方案以創(chuàng)建小批量。

然后，制作小批量的程序如下：

from torch_geometric.loader import ClusterData, ClusterLoadertorch.manual_seed(12345)
cluster_data = ClusterData(data, num_parts=128)  # 1. Create subgraphs.
train_loader = ClusterLoader(cluster_data, batch_size=32, shuffle=True)  # 2. Stochastic partioning scheme.print()
total_num_nodes = 0
for step, sub_data in enumerate(train_loader):print(f'Step {step + 1}:')print('=======')print(f'Number of nodes in the current batch: {sub_data.num_nodes}')print(sub_data)print()total_num_nodes += sub_data.num_nodesprint(f'Iterated over {total_num_nodes} of {data.num_nodes} nodes!')

???????在這里，我們將初始圖劃分為128個(gè)分區(qū)，并使用32個(gè)子圖的batch_size來(lái)形成mini-batches（每個(gè)epoch4批）。正如我們所看到的，在一個(gè)epoch之后，每個(gè)節(jié)點(diǎn)都被精確地看到了一次。

???????Cluster-GCN的偉大之處在于它不會(huì)使GNN模型的實(shí)現(xiàn)復(fù)雜化。我們構(gòu)造如下一個(gè)簡(jiǎn)單模型：

???????這種圖神經(jīng)網(wǎng)絡(luò)的訓(xùn)練與用于圖分類任務(wù)的圖神經(jīng)網(wǎng)絡(luò)訓(xùn)練非常相似。我們現(xiàn)在不再以全批處理的方式對(duì)圖進(jìn)行操作，而是對(duì)每個(gè)小批進(jìn)行迭代，并相互獨(dú)立地優(yōu)化每個(gè)批：

model = GCN(hidden_channels=16)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
criterion = torch.nn.CrossEntropyLoss()def train():model.train()for sub_data in train_loader:  # Iterate over each mini-batch.out = model(sub_data.x, sub_data.edge_index)  # Perform a single forward pass.loss = criterion(out[sub_data.train_mask], sub_data.y[sub_data.train_mask])  # Compute the loss solely based on the training nodes.loss.backward()  # Derive gradients.optimizer.step()  # Update parameters based on gradients.optimizer.zero_grad()  # Clear gradients.def test():model.eval()out = model(data.x, data.edge_index)pred = out.argmax(dim=1)  # Use the class with highest probability.accs = []for mask in [data.train_mask, data.val_mask, data.test_mask]:correct = pred[mask] == data.y[mask]  # Check against ground-truth labels.accs.append(int(correct.sum()) / int(mask.sum()))  # Derive ratio of correct predictions.return accsfor epoch in range(1, 51):loss = train()train_acc, val_acc, test_acc = test()print(f'Epoch: {epoch:03d}, Train: {train_acc:.4f}, Val Acc: {val_acc:.4f}, Test Acc: {test_acc:.4f}')

???????在本文中，我們向您介紹了一種將 scale GNNs to large graphs的方法，否則這些圖將不適合GPU內(nèi)存。

本文內(nèi)容參考：PyG官網(wǎng)