怎么使用Pytorch+PyG實現GAT

本篇內容主要講解“怎么使用Pytorch+PyG實現GAT”，感興趣的朋友不妨來看看。本文介紹的方法操作簡單快捷，實用性強。下面就讓小編來帶大家學習“怎么使用Pytorch+PyG實現GAT”吧!

導入庫和數據

GAT（圖注意力網絡）是常見的圖神經網絡結構之一，它使用注意力機制來對節點進行特征加權，并考慮其鄰居節點的交互。

首先，我們需要導入PyTorch和PyG庫，然后準備好我們的數據。例如，我們可以使用以下方式生成一個簡單的隨機數據集：

from torch_geometric.datasets import Planetoid
dataset = Planetoid(root='/tmp/Cora', name='Cora')
train_loader = DataLoader(dataset[0], batch_size=128, shuffle=True)
test_loader = DataLoader(dataset[0], batch_size=128, shuffle=False)

其中， Planetoid 是PyG提供的圖形數據集之一。這里我們選擇了 Cora 數據集并存儲到 /tmp/Cora 文件夾中。然后我們將該數據集分成訓練集和測試集，設置相應的加載器。

定義模型結構

接下來，我們需要定義GAT模型的結構。通過PyTorch和PyG，我們可以自己定義完整的GAT模型或者利用現有的庫函數快速構建模型。在這里，我們將使用 torch_geometric.nn.GATConv 函數逐層堆疊多個圖注意力層來實現GAT模型。以下是GAT模型定義的示例代碼：

import torch.nn.functional as F
from torch_geometric.nn import GATConv
class Net(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Net, self).__init__()
        self.num_layers = 2
        self.conv1 = GATConv(in_channels=in_channels, out_channels=16, heads=8, dropout=0.6)
        self.conv2 = GATConv(in_channels=16*8, out_channels=out_channels, heads=1, concat=False, dropout=0.6)
    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = F.dropout(x, p=0.6, training=self.training)
        x = F.elu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.6, training=self.training)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

上述代碼中，我們定義了一個 Net 類用于構建GAT網絡，接收輸入通道數和輸出通道數作為參數。例如，我們可以按照以下方式創建一個將 CORD 參量作為輸入特征向量大小、64 個隱藏節點（每個注意力頭）。并將數字類別作為輸出大小的GAT模型：

model = Net(in_channels=dataset.num_features, out_channels=dataset.num_classes)

其中 num_features 和 num_classes 是PyG數據集中包含的屬性。

定義訓練函數

然后，我們需要定義訓練函數來訓練我們的GAT神經網絡。在這里，我們將使用交叉熵損失和Adam優化器進行訓練，并在每一個epoch結束時計算準確率并打印出來。以下是訓練函數的示例代碼：

import torch.optim as optim
from tqdm import tqdm
def train(model, loader, optimizer, loss_fn):
    model.train()
    correct = 0
    total_loss = 0
    for data in tqdm(loader, desc='Training'):
        optimizer.zero_grad()
        out = model(data)
        pred = out.argmax(dim=1)
        loss = loss_fn(out[data.train_mask], data.y[data.train_mask])
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.num_graphs
        correct += pred[data.train_mask].eq(data.y[data.train_mask]).sum().item()
    return total_loss / len(loader.dataset), correct / len(data.train_mask)

在上述代碼中，我們遍歷加載器中的每個數據批次，并對模型進行培訓。對于每個圖數據批次，我們計算網絡輸出、預測和損失，然后通過反向傳播來更新權重。最后，我們將總損失和正確率記錄下來并返回。

定義測試函數

接下來，我們還需要定義測試函數來測試我們的GAT神經網絡性能表現。我們將利用與訓練函數相同的輸出參數進行測試，并打印出最終的測試準確率。以下是測試函數的示例代碼：

def test(model, loader, loss_fn):
    model.eval()
    correct = 0
    total_loss = 0
    with torch.no_grad():
        for data in tqdm(loader, desc='Testing'):
            out = model(data)
            pred = out.argmax(dim=1)
            loss = loss_fn(out[data.test_mask], data.y[data.test_mask])
            total_loss += loss.item() * data.num_graphs
            correct += pred[data.test_mask].eq(data.y[data.test_mask]).sum().item()
    return total_loss / len(loader.dataset), correct / len(data.test_mask)

在上述代碼中，我們對測試數據集中的所有數據進行了循環，并計算網絡的輸出和預測。我們記錄下總損失和正確分類的數據量，并返回損失和準確率之間的比率。

訓練模型并評估訓練結果

最后，我們可以使用前面定義過的函數來定義主函數，從而完成GAT神經網絡的訓練和測試。以下是主函數的示例代碼：

if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = Net(in_channels=dataset.num_features, out_channels=dataset.num_classes).to(device)
    train_loader = DataLoader(dataset[0], batch_size=128, shuffle=True)
    test_loader = DataLoader(dataset[0], batch_size=128, shuffle=False)
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    loss_fn = nn.CrossEntropyLoss()
    for epoch in range(1, 201):
        train_loss, train_acc = train(model, train_loader, optimizer, loss_fn)
        test_loss, test_acc = test(model, test_loader, loss_fn)
        print(f'Epoch {epoch:03d}, Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, '
              f'Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f}')

通過上述代碼，我們就可以完成GAT神經網絡的訓練和測試。我們使用 DataLoader 函數進行數據加載，設置學習率、損失函數、訓練輪數等超參數。最后，我們可以在屏幕上看到每個時代的準確率和損失值，并通過它們評估模型的訓練表現。

到此，相信大家對“怎么使用Pytorch+PyG實現GAT”有了更深的了解，不妨來實際操作一番吧！這里是億速云網站，更多相關內容可以進入相關頻道進行查詢，關注我們，繼續學習！