图学习——03.PyG的使用

PYG官网：https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html

Installtion

PyTorch 1.12.* and CUDA 11.6, type:

pip install torch-scatter -f https://data.pyg.org/whl/torch-1.12.0+cu116.html
pip install torch-sparse -f https://data.pyg.org/whl/torch-1.12.0+cu116.html
pip install torch-geometric

PyTorch 1.11.* and CUDA 11.3, type:

pip install torch-scatter -f https://data.pyg.org/whl/torch-1.11.0+cu113.html
pip install torch-sparse -f https://data.pyg.org/whl/torch-1.11.0+cu113.html
pip install torch-geometric

Introduction

官网详细举例
下面是对举例的解释和注释，记录学习内容

一些方法的使用

# -*- coding = utf-8 -*-
# @Time : 2022/10/16 10:18
# @Author : 头发没了还会再长
# @File : pyg.py
# @Software : PyCharm

import torch
from torch_geometric.data import Data

# 边（0,1）（1,0）（1,2）（2,1）
edge_index = torch.tensor([[0, 1, 1, 2],
                           [1, 0, 2, 1]], dtype=torch.long)
# 每一个节点的特征值 一维的向量表示
x = torch.tensor([[-1], [0], [1]], dtype=torch.float)
# 将边edge_index和节点x 传到Data类里 得到data数据集
data = Data(x=x, edge_index=edge_index)

# 相关api的使用
print(format(data)) # Data(x=[3, 1], edge_index=[2, 4]) [3, 1]表示三个节点，每个节点的特征维度是1，[2, 4]表示4/2条边
print(data.keys) # ['edge_index', 'x']
print(data['x']) # 通过关键字得到value tensor([[-1.],[ 0.],[ 1.]])
print('edge_attr' in data) # False 判断边是否有属性
print(data.num_nodes) # 3 节点数量
print(data.num_edges) # 4 边数量
print(data.num_node_features) # 1 节点属性的shape
print(data.contains_isolated_nodes()) # Fasle 是否有孤立的节点
print(data.contains_self_loops()) # False 是否形成环
print(data.is_directed()) # False 是否是有向图

# Transfer data object to GPU.
device = torch.device('cuda')
data = data.to(device)

数据集TUDataset

# -*- coding = utf-8 -*-
# @Time : 2022/10/16 11:30
# @Author : 头发没了还会再长
# @File : Dataset.py
# @Software : PyCharm

from torch_geometric.datasets import TUDataset

# 数据集 TUDataset
dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES')
print(len(dataset)) # 600 数据集的大小
print(dataset.num_classes) # 6 数据集的类别
print(dataset.num_node_features) # 3 节点的特征数

# 取数据集里的第一个图
data = dataset[0]
print(data)# Data(edge_index=[2, 168], x=[37, 3], y=[1]) 37个有3个特征值的节点x，168/2条边 类别为1
print(data.is_directed()) #False 无向图

# 将数据集切片 分成训练集和测试集，比例为9:1
train_dataset = dataset[:540]
test_dataset = dataset[540:]
dataset = dataset.shuffle()

# mini_batches PyG有自己的DataLoader
dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES', use_node_attr=True)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

for batch in loader:
    print(batch) # DataBatch(edge_index=[2, 3672], x=[957, 21], y=[32], batch=[957], ptr=[33])
    print(batch.num_graphs) # 32

数据集Cora

from torch_geometric.datasets import Planetoid

# 数据集Cora
dataset = Planetoid(root='/tmp/Cora', name='Cora')
print(len(dataset)) # 1 只有一个图
print(dataset.num_classes()) # 7 节点类别数7
print(dataset.num_node_features) # 1443 节点特征数1443

# 取数据集里的第一个图
data = dataset[0]
print(data) # Data(edge_index=[2, 10556], test_mask=[2708],
         # train_mask=[2708], val_mask=[2708], x=[2708, 1433], y=[2708])
         # train_mask 表示一共2708个节点，当是训练集时，节点对应的特征值是1，否则就是0 val和train类似

# 每一类数据集节点掩蔽的数量
data.train_mask.sum().item() # 140
data.val_mask.sum().item() #500
data.test_mask.sum().item() # 1000

一个两次GCN的模型举例

# -*- coding = utf-8 -*-
# @Time : 2022/10/16 12:27
# @Author : 头发没了还会再长
# @File : GCN.py
# @Software : PyCharm
from torch_geometric.datasets import Planetoid
import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

dataset = Planetoid(root='/tmp/Cora', name='Cora')

# implement a two-layer GCN
class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_node_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)

        return F.log_softmax(x, dim=1)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GCN().to(device)
data = dataset[0].to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

model.train()
for epoch in range(200):
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

model.eval()
pred = model(data).argmax(dim=1)
correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()
acc = int(correct) / int(data.test_mask.sum())
print(f'Accuracy: {acc:.4f}')
# Accuracy: 0.8150