torch-geometric包的使用

1.torch-geometric的安装

1.1安装注意点

        1.需要安装他的依赖包

        2.搭配使用的torch版本>=1.8.1，最好直接安装torch==1.10.1

        3.需要将依赖包直接下载到本地（GitHub - pyg-team/pytorch_geometric: Graph Neural Network Library for PyTorch
 ）点击进入here的位置

​

         注意依赖包需要对应版本的python以及torch

        4.安装时直接pip install "依赖项的完整路径"

        5.最后pip install torch-geometric

        GitHub - pyg-team/pytorch_geometric: Graph Neural Network Library for PyTorch

2.内置数据集

2.1数据集：Zachary's karate club network.

​

from torch_geometric.datasets import KarateClub

dataset = KarateClub()
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}')
##结果如下

Dataset: KarateClub():
======================
Number of graphs: 1
Number of features: 34
Number of classes: 4

 2.1.1查看图的信息

        查看该数据第一张图（一共一张图）的信息

        

data = dataset[0]  # Get the first graph object.

print(data)
##结果如下
Data(x=[34, 34], edge_index=[2, 156], y=[34], train_mask=[34])

 

edge_index = data.edge_index
print(edge_index.t())
#结果如下
tensor([[ 0,  1],
        [ 0,  2],
        [ 0,  3],
        [ 0,  4],
        [ 0,  5],
        [ 0,  6],
        [ 0,  7],
        [ 0,  8],
        ...

2.1.2图数据的可视化

from torch_geometric.utils import to_networkx

G = to_networkx(data, to_undirected=True)
visualize_graph(G, color=data.y)

 3.GCN的构建

3.1构建GCN模型

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


class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        torch.manual_seed(1234)
        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)

#结果如下
GCN(
  (conv1): GCNConv(34, 4)
  (conv2): GCNConv(4, 4)
  (conv3): GCNConv(4, 2)
  (classifier): Linear(in_features=2, out_features=4, bias=True)
)

3.2训练模型

import time

model = GCN()
criterion = torch.nn.CrossEntropyLoss()  # Define loss criterion.
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)  # Define optimizer.

def train(data):
    optimizer.zero_grad()  
    out, h = model(data.x, data.edge_index) #h是两维向量，主要是为了咱们画个图 
    loss = criterion(out[data.train_mask], data.y[data.train_mask])  # semi-supervised
    loss.backward()  
    optimizer.step()  
    return loss, h

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

本文参考链接