PyG官方文檔文章來源地址http://www.zghlxwxcb.cn/news/detail-667285.html

# Install torch geometric
!pip install -q torch-scatter -f https://pytorch-geometric.com/whl/torch-1.10.2+cu102.html
!pip install -q torch-sparse -f https://pytorch-geometric.com/whl/torch-1.10.2+cu102.html
!pip install -q torch-geometric

import torch
import networkx as nx
import matplotlib.pyplot as plt

1.內置數據集（以KarateClub為例）

from torch_geometric.datasets import KarateClub

dataset = KarateClub()
print(f'Dataset: {dataset}:')
print('======================')
# 圖的數量
print(f'Number of graphs: {len(dataset)}')
# 每個節(jié)點的特征尺寸
print(f'Number of features: {dataset.num_features}')
# 節(jié)點的類別數量
print(f'Number of classes: {dataset.num_classes}')

# 獲取具體的圖
data = dataset[0]
print(data)
print('==============================================================')

# 獲取圖的屬性
print(f'Number of nodes: {data.num_nodes}')
print(f'Number of edges: {data.num_edges}')
print(f'Average node degree: {(2*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:.2f}')
print(f'Contains isolated nodes: {data.has_isolated_nodes()}')
print(f'Contains self-loops: {data.has_self_loops()}')
print(f'Is undirected: {data.is_undirected()}')

# 取出的圖的數據對象為Data類型，包含以下屬性
# 1. edge_index 每條邊的兩個端點的索引組成的元組
# 2. x 節(jié)點特征[節(jié)點數量，特征維數]
# 3. y 節(jié)點標簽（類別），每個節(jié)點只分配一個類別
# 4. train_mask 
Data(edge_index=[2, 156], x=[34, 34], y=[34], train_mask=[34])
print(data)

# 獲取所有的邊
print(data.edge_idx.T)

2.可視化

def visualize(h, color, epoch=None, loss=None, accuracy=None):
    plt.figure(figsize=(7,7))
    plt.xticks([])
    plt.yticks([])
    
    if torch.is_tensor(h):
        h = h.detach().cpu().numpy()
        plt.scatter(h[:, 0], h[:, 1], s=140, c=color, cmap="Set2")
        if epoch is not None and loss is not None and accuracy['train'] is not None and accuracy['val'] is not None:
            plt.xlabel((f'Epoch: {epoch}, Loss: {loss.item():.4f} \n'
                       f'Training Accuracy: {accuracy["train"]*100:.2f}% \n'
                       f' Validation Accuracy: {accuracy["val"]*100:.2f}%'),
                       fontsize=16)
    else:
        # networkx的draw_networkx
        nx.draw_networkx(h, pos=nx.spring_layout(h, seed=42), with_labels=False, node_color=color, cmap="Set2")   
    
    plt.show()

from torch_geometric.utils import to_networkx
# 將Data類型轉換成networkx
G = to_networkx(data, to_undirected=True)
# 將圖可視化，節(jié)點顏色為節(jié)點的類型
visualize(G, color=data.y)

3.搭建GNN(以GCN為例)

import torch
from torch.nn import Linear
from torch_geometric.nn import GCNConv

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_features, 4)
        self.conv2 = GCNConv(4, 4)
        self.conv3 = GCNConv(4, 2)
        self.classifier = Linear(2, dataset.num_classes)
    def forward(self, x, edge_index):
        h = self.conv1(x, edge_index)
        h = h.tanh()
        h = self.conv2(h, edge_index)
        h = h.tanh()
        h = self.conv3(h, edge_index)
        h = h.tanh()
        out = self.classifier(h)
        return out, h

model = GCN()
print(model)

# 節(jié)點分類
model = GCN()

out, h = model(data.x, data.edge_index)
print(f'Embedding shape: {list(h.shape)}')

visualize(h, color=data.y)

4.在KarateClub數據集上訓練

import time
model = GCN()

# 交叉熵損失，Adam優(yōu)化器
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

def train(data):
    optimizer.zero_grad()
    out, h  = model(data.x, data.edge_index)
    # 只對train_mask的節(jié)點計算loss
    loss = criterion(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    
    accuracy = {}
    # torch.argmax 取置信度最大的一類
    predicted_classes = torch.argmax(out[data.train_mask], axis=1) # [0.6, 0.2, 0.7, 0.1] -> 2
    target_classes = data.y[data.train_mask]
    accuracy['train'] = torch.mean(torch.where(predicted_classes == target_classes, 1, 0).float())
    
    predicted_classes = torch.argmax(out, axis=1)
    target_classes = data.y
    accuracy['val'] = torch.mean(torch.where(predicted_classes == target_classes, 1, 0).float())
    
    return loss, h, accuracy

for epoch in range(500):
    loss, h, accuracy = train(data)
    if epoch % 10 == 0:
        visualize(h, color=data.y, epoch=epoch, loss=loss, accuracy=accuracy)
        time.sleep(0.3)

到了這里，關于PyTorch Geometric基本教程的文章就介紹完了。如果您還想了解更多內容，請在右上角搜索TOY模板網以前的文章或繼續(xù)瀏覽下面的相關文章，希望大家以后多多支持TOY模板網！