图神经网络框架-PyTorch Geometric(PyG)的使用及踩坑

import torch
from torch_geometric.data import Data

# 边的连接信息
# 注意，无向图的边要定义两次
edge_index = torch.tensor(
    [
        # 这里表示节点0和1有连接，因为是无向图
        # 那么1和0也有连接
        # 上下对应着看
        [0, 1, 1, 2],
        [1, 0, 2, 1],
    ],
    # 指定数据类型
    dtype=torch.long
)
# 节点的属性信息
x = torch.tensor(
    [
        # 三个节点
        # 每个节点的属性向量维度为1
        [-1],
        [0],
        [1],
    ]
)
# 实例化为一个图结构的数据
data = Data(x=x, edge_index=edge_index)
# 查看图数据
print(data)
# 图数据中包含什么信息
print(data.keys)
# 查看节点的属性信息
print(data['x'])
# 节点数
print(data.num_nodes)
# 边数
print(data.num_edges)
# 节点属性向量的维度
print(data.num_node_features)
# 图中是否有孤立节点
print(data.has_isolated_nodes())
# 图中是否有环
print(data.has_self_loops())
# 是否是有向图
print(data.is_directed())

from torch_geometric.datasets import TUDataset

# 导入数据集
dataset = TUDataset(
    # 指定数据集的存储位置
    # 如果指定位置没有相应的数据集
    # PyG会自动下载
    root='../data/ENZYMES',
    # 要使用的数据集
    name='ENZYMES',
)
# 数据集的长度
print(len(dataset))
# 数据集的类别数
print(dataset.num_classes)
# 数据集中节点属性向量的维度
print(dataset.num_node_features)
# 600个图，我们可以根据索引选择要使用哪个图
data = dataset[0]
print(data)
# 随机打乱数据集
dataset = dataset.shuffle()

from torch_geometric.loader import DataLoader
from torch_geometric.datasets import TUDataset


# 数据集
dataset = TUDataset(
    root='../data/ENZYMES',
    name='ENZYMES',
    use_node_attr=True,
)
# 建立数据集加载器
# 每次加载32个数据到内存中
loader = DataLoader(
    # 要加载的数据集
    dataset=dataset,
    # ENZYMES包含600个图
    # 每次加载32个
    batch_size=32,
    # 每次加入进来之后是否随机打乱数据（可以增加模型的泛化性）
    shuffle=True
)
for batch in loader:
    print(batch)
    print(batch.num_graphs)

from abc import ABC

import torch
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree


# 定义GCN空域图卷积神经网络
class GCNConv(MessagePassing, ABC):
    # 网络初始化
    def __init__(self, in_channels, out_channels):
        """
        :param in_channels: 节点属性向量的维度
        :param out_channels: 经过图卷积之后，节点的特征表示维度
        """
        # 定义伽马函数为求和函数,aggr='add'
        super(GCNConv, self).__init__(aggr='add')
        # 定义最里面那个线性变换
        # 具体到实现中就是一个线性层
        self.linear_change = torch.nn.Linear(in_channels, out_channels)

    # 定义信息汇聚函数
    def message(self, x_j, norm):
        # 正则化
        # norm.view(-1,1)将norm变为一个列向量
        # x_j是节点的特征表示矩阵
        return norm.view(-1, 1) * x_j

    # 前向传递，进行图卷积
    def forward(self, x, edge_index):
        """
        :param x:图中的节点，维度为[节点数,节点属性相邻维度数]
        :param edge_index: 图中边的连接信息,维度为[2,边数]
        :return:
        """
        # 添加节点到自身的环
        # 因为节点最后面汇聚相邻节点信息时包含自身
        # add_self_loops会在edge_index边的连接信息表中，
        # 添加形如[i,i]这样的信息
        # 表示一个节点到自身的环
        # 函数返回[边的连接信息，边上的属性信息]
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        # 进行线性变换
        x = self.linear_change(x)
        # 计算外面的正则化
        row, col = edge_index
        # 获取节点的度
        deg = degree(col, x.size(0), dtype=x.dtype)
        # 带入外面的正则化公式
        deg_inv_sqrt = deg.pow(-0.5)
        # 将未知的值设为0，避免下面计算出错
        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
        # 正则化部分
        norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]
        # 进行信息传递和融合
        # propagate会自动调用self.message函数，并将参数传递给它
        return self.propagate(edge_index, x=x, norm=norm)


# 测试我们刚才定义的图卷积神经网络
if __name__ == '__main__':
    # 实例化一个图卷积神经网络
    # 并假设图节点属性向量的维度为16，图卷积出来的节点特征表示向量维度为32
    conv = GCNConv(16, 32)
    # 随机生成一个节点属性向量
    # 5个节点，属性向量为16维
    x = torch.randn(5, 16)
    # 随机生成边的连接信息
    # 假设有3条边
    edge_index = [
        [0, 1, 1, 2, 1, 3],
        [1, 0, 2, 1, 3, 1]
    ]
    edge_index = torch.tensor(edge_index, dtype=torch.long)
    # 进行图卷积
    output = conv(x, edge_index)
    # 输出卷积之后的特征表示矩阵
    print(output.data)

import torch
from torch.nn import Sequential as Seq
from torch.nn import Linear, ReLU
from torch_geometric.nn import MessagePassing


# 定义EdgeConv图卷积神经网络
class EdgeConv(MessagePassing):
    # 初始化图卷积神经网络
    def __init__(self, in_channels, out_channels):
        # 定义伽马函数为求最大值函数
        super().__init__(aggr='max')
        # 定义一个前馈神经网络
        self.mlp = Seq(
            # 线性层,后面信息汇聚函数之后的输入是2*in_channels
            Linear(2 * in_channels, out_channels),
            # 激活函数
            ReLU(),
            # 输出层
            Linear(out_channels, out_channels)
        )

    # 定义信息汇聚函数
    def message(self, x_i, x_j):
        tmp = torch.cat([x_i, x_j - x_i], dim=1)
        # cat之后tmp的维度为[边数,2*in_channels]
        return self.mlp(tmp)

    # 前向传递，进行图卷积
    def forward(self, x, edge_index):
        # x是节点属性向量矩阵
        # edge_index是边的连接信息
        # 进行信息的传递、融合
        return self.propagate(edge_index, x=x)

import torch
from torch_geometric.data import InMemoryDataset, download_url


# 实现In Memory Dataset的通用模板
class MyDataset(InMemoryDataset):
    # 初始化
    def __init__(self, root, transfrom=None, pre_transform=None):
        # root是数据集的根目录
        super(MyDataset, self).__init__(root, transfrom, pre_transform)
        # 加载数据集
        self.data, self.slices = torch.load(self.processed_paths[0])

    def raw_file_names(self) -> Union[str, List[str], Tuple]:
        return ['file_1', 'file_2', ...]

    def processed_file_names(self) -> Union[str, List[str], Tuple]:
        return ['data.pt']

    def download(self):
        # 将数据集下载到raw_dir文件夹中
        download_url(url, self.raw_dir)

    def process(self):
        data_list = [...]
        # 进行数据过滤
        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]
        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]
        # self.collate将所有数据组合在一起,加速存储
        # data是组合之后的数据
        # slices是分割方式，告诉PyG如何将data还原为原先的数据
        data, slices = self.collate(data_list)
        # 保存数据
        torch.save((data, slices), self.processed_paths[0])

import os.path as osp
import torch
from torch_geometric.data import Dataset, download_url


class MyDataset(Dataset):
    # 初始化
    def __init__(self, root, transform=None, pre_transform=None):
        super(MyDataset, self).__init__(root, transform, pre_transform)

    def raw_file_names(self) -> Union[str, List[str], Tuple]:
        return ['file_1', 'file_2', ...]

    def processed_file_names(self) -> Union[str, List[str], Tuple]:
        return ['data_1.pt', ...]

    def download(self):
        path = download_url(url, self.raw_dir)

    def process(self):
        i = 0
        for raw_path in self.raw_paths:
            # 读取数据
            data = Data(...)
            # 过滤数据集
            if self.pre_filter is not None and not self.pre_filter(data):
                pass
            if self.pre_transform is not None:
                data = self.pre_transform(data)
            # 保存数据
            torch.save(data, osp.join(self.processed_dir, 'data_{}.pt'.format(i)))
            i += 1

    def len(self):
        return len(self.processed_file_names)

    def get(self,idx):
        data = torch.load(osp.join(self.processed_dir, 'data_{}.pt'.format(idx)))
        return data

from typing import Any

import torch
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader


# 定义图数据
class PairData(Data):
    def __init__(self, edge_index_s=None, x_s=None, edge_index_t=None, x_t=None):
        # 每个数据中包含两个图s,t
        """
        :param edge_index_s: 图s的连接关系
        :param x_s: 图s的节点属性矩阵
        :param edge_index_t: 图t的连接关系
        :param x_t: 图t的节点属性矩阵
        """
        super(PairData, self).__init__()
        self.edge_index_s = edge_index_s
        self.x_s = x_s
        self.edge_index_t = edge_index_t
        self.x_t = x_t

    def __inc__(self, key: str, value: Any, *args, **kwargs) -> Any:
        # 如果要合并的是图s
        # 那么告诉PyG图s的节点数
        if key == 'edge_index_s':
            return self.x_s.size(0)
        # 如果要合并的是图t
        # 那么告诉PyG图t的节点数
        if key == 'edge_index_t':
            return self.x_t.size(0)
        # 其它情况默认
        else:
            return super().__inc__(key, value, *args, **kwargs)


# 下面验证一下我们上面定义的合并方法
# 定义图s
edge_index_s = torch.tensor([
    [0, 0, 0, 0],
    [1, 2, 3, 4],
])
x_s = torch.randn(5, 16)
# 定义图t
edge_index_t = torch.tensor([
    [0, 0, 0],
    [1, 2, 3],
])
x_t = torch.randn(4, 16)  # 4 nodes.
# 验证，简单定义数据集包含两个数据
data = PairData(edge_index_s, x_s, edge_index_t, x_t)
data_list = [data, data]
# batch_size=2
# follow_batch描述节点信息
loader = DataLoader(data_list, batch_size=2, follow_batch=['x_s', 'x_t'])
# 验证PyG是否按照我们定义的方式有效的合并了一个batch的数据
batch = next(iter(loader))
# 查看合并为一个batch的数据
print(batch)
# 查看batch中的s(这个是两个原数据中s的组合，作为一个)
print(batch.edge_index_s)
# 查看batch中的t
print(batch.edge_index_t)

import torch
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader


# 定义二分图结构
class BipartiteData(Data):
    def __init__(self, edge_index=None, x_s=None, x_t=None):
        super().__init__()
        # 包含一组边
        # 两组节点
        self.edge_index = edge_index
        self.x_s = x_s
        self.x_t = x_t

    # 定义每个batch的合并方式
    def __inc__(self, key, value, *args, **kwargs):
        # 如果要合并两个图的边连接信息
        if key == 'edge_index':
            # 左边（边连接信息的第一行）按照第一组节点数合并
            # 右边（边连接信息的第二行）按照第二组节点数合并
            return torch.tensor([[self.x_s.size(0)], [self.x_t.size(0)]])
        else:
            return super().__inc__(key, value, *args, **kwargs)


edge_index = torch.tensor([
    [0, 0, 1, 1],
    [0, 1, 1, 2],
])
x_s = torch.randn(2, 16)
x_t = torch.randn(3, 16)
data = BipartiteData(edge_index, x_s, x_t)
data_list = [data, data]
loader = DataLoader(data_list, batch_size=2)
batch = next(iter(loader))
print(batch)
print(batch.edge_index)

import os.path as osp

import torch
import pandas as pd
from sentence_transformers import SentenceTransformer

from torch_geometric.data import HeteroData, download_url, extract_zip
from torch_geometric.transforms import ToUndirected, RandomLinkSplit

# 数据集下载地址
url = 'https://files.grouplens.org/datasets/movielens/ml-latest-small.zip'
# 数据集存放路径
root = osp.join(osp.dirname(osp.realpath(__file__)), '../data/MovieLens')
# 下载数据集并进行解压
extract_zip(download_url(url, root), root)
# 获取movies.csv，ratings.csv文件
movie_path = osp.join(root, 'ml-latest-small', 'movies.csv')
rating_path = osp.join(root, 'ml-latest-small', 'ratings.csv')
# 利用pandas查看数据集
print(pd.read_csv(movie_path).head())
print(pd.read_csv(rating_path).head())


# 将电影名那列
# 利用嵌入模型将每个电影名用向量表示(Embedding)
class SequenceEncoder(object):
    # 初始化
    # 指定我们使用的嵌入模型
    # 和使用的设备
    def __init__(self, model_name='all-MiniLM-L6-v2', device=None):
        # 使用的设备
        self.device = device
        # 使用的嵌入模型名
        self.model = SentenceTransformer(model_name, device=device)

    # 嵌入模型不参与后续图神经网络的训练
    @torch.no_grad()
    def __call__(self, df):
        x = self.model.encode(
            # 要进行嵌入的值
            df.values,
            # 显示处理进度
            show_progress_bar=True,
            # 转换为PyTorch的张量
            convert_to_tensor=True,
            # 使用的设备
            device=self.device
        )
        return x.cpu()


# 将电影类型那列进行嵌入表示
class GenresEncoder(object):

    # 分隔符
    def __init__(self, sep='|'):
        self.sep = sep

    def __call__(self, df):
        # 分割出所有的电影类型
        # 后面两个for的逻辑是：
        # for col in df.values取出每一行的值
        # for g in col.split(self.sep)将取出来的值用指定的分隔符进行分割
        # set(g)将分割之后的结果转换为集合,去重
        genres = set(g for col in df.values for g in col.split(self.sep))
        # 将电影类型用数字表示
        mapping = {genre: i for i, genre in enumerate(genres)}
        # 用multi-hot形式表示电影的类型
        x = torch.zeros(len(df), len(mapping))
        for i, col in enumerate(df.values):
            for genre in col.split(self.sep):
                x[i, mapping[genre]] = 1
        return x


# 从CSV文件中读取信息，建立二分图中节点的信息
def load_node_csv(path, index_col, encoders=None, **kwargs):
    """
    :param path: CSV文件路径
    :param index_col: 文件中的索引列，也就是节点所在的列
    :param encoders:节点嵌入器
    :param kwargs:
    :return:
    """
    df = pd.read_csv(path, index_col=index_col, **kwargs)
    # 将索引用数字表示
    mapping = {index: i for i, index in enumerate(df.index.unique())}
    # 节点属性向量矩阵
    x = None
    # 如果嵌入器非空
    if encoders is not None:
        # 对相应的列进行嵌入
        # 获取嵌入向量表示
        xs = [encoder(df[col]) for col, encoder in encoders.items()]
        x = torch.cat(xs, dim=-1)

    return x, mapping


# 获取节点信息
# 处理movies.csv表，将'电影名','电影类型'列转换为嵌入向量的表示形式
movie_x, movie_mapping = load_node_csv(
    movie_path, index_col='movieId', encoders={
        # 电影名列的嵌入器
        'title': SequenceEncoder(),
        # 电影类型列的嵌入器
        'genres': GenresEncoder()
    })
# 处理ratings.csv表,将用户ID用PyTorch中的张量表示
user_x, user_mapping = load_node_csv(rating_path, index_col='userId')
# 建立异质图（这里具体是一个二分图）
# HeteroData()是PyG中内置的一个表示异质图的数据结构
data = HeteroData()
# 加入不同类型节点的信息
# 加入用户信息，用户没有属性向量
# 只需要告诉PyG有多少个用户节点就可以
data['user'].num_nodes = len(user_mapping)
# 告诉PyG 电影的属性向量矩阵，PyG会根据x推断出电影节点的个数
data['movie'].x = movie_x
print(data)


# 建立用户和电影之间边的信息
# 将用户对电影的评分转换为PyTorch中的张量
# 方便后续模型的训练
class IdentityEncoder(object):

    def __init__(self, dtype=None):
        self.dtype = dtype

    def __call__(self, df):
        return torch.from_numpy(df.values).view(-1, 1).to(self.dtype)


# 建立二分图边的连接信息
def load_edge_csv(path, src_index_col, src_mapping, dst_index_col, dst_mapping,
                  encoders=None, **kwargs):
    """
    :param path: CSV表的路径
    :param src_index_col: 二分图左边节点来源于CSV表的哪一列，比如'user_id'这列
    :param src_mapping:将user_id映射为节点编号，我们前面定义的user_mapping
    :param dst_index_col:同理，二分图右边电影节点
    :param dst_mapping:
    :param encoders:边的嵌入器
    :param kwargs:
    :return:
    """
    df = pd.read_csv(path, **kwargs)
    # 建立连接信息
    src = [src_mapping[index] for index in df[src_index_col]]
    dst = [dst_mapping[index] for index in df[dst_index_col]]
    # 注意这里edge_index维度为[2,边数]
    edge_index = torch.tensor([src, dst])
    # 边的属性信息
    edge_attr = None
    # 如果嵌入器非空
    if encoders is not None:
        edge_attrs = [encoder(df[col]) for col, encoder in encoders.items()]
        edge_attr = torch.cat(edge_attrs, dim=-1)

    return edge_index, edge_attr


# 获取二分图边的信息
edge_index, edge_label = load_edge_csv(
    rating_path,
    # 二分图左边是用户
    src_index_col='userId',
    src_mapping=user_mapping,
    # 右边是电影
    dst_index_col='movieId',
    dst_mapping=movie_mapping,
    encoders={'rating': IdentityEncoder(dtype=torch.long)},
)
# 将二分图中的边命名为('user', 'rates', 'movie')
data['user', 'rates', 'movie'].edge_index = edge_index
data['user', 'rates', 'movie'].edge_label = edge_label
print(data)

# 到此我们的异质图(这里是一个二分图)数据集就构建完毕了
# 下面进一步将其转换为一个真正可以进行训练的数据集
# 转换为无向图
data = ToUndirected()(data)
# 删除相反方向边的属性信息，因为没有电影对用户的评分数据
del data['movie', 'rev_rates', 'user'].edge_label

# 按照一定比例分割数据集为训练集、测试集、验证集
transform = RandomLinkSplit(
    num_val=0.05,
    num_test=0.1,
    # 负采样比率
    # 不用负采样，全部输入进行训练
    neg_sampling_ratio=0.0,
    # 告诉PyG边的连接关系
    edge_types=[('user', 'rates', 'movie')],
    rev_edge_types=[('movie', 'rev_rates', 'user')],
)
# 分割数据集
train_data, val_data, test_data = transform(data)
print(train_data)
print(val_data)
print(test_data)

from torch_geometric.data import HeteroData

# HeteroData是PyG自带的一个异质图数据结构
data = HeteroData()
# 添加节点的信息
data['paper'].x = ...
data['author'].x = ...
data['institution'].x = ...
data['field_of_study'].x = ...
# 添加边的连接信息
data['paper', 'cites', 'paper'].edge_index = ...  
data['author', 'writes', 'paper'].edge_index = ...  
data['author', 'affiliated_with', 'institution'].edge_index = ...  
data['author', 'has_topic', 'institution'].edge_index = ...  
# 添加边的属性信息
data['paper', 'cites', 'paper'].edge_attr = ...  
data['author', 'writes', 'paper'].edge_attr = ...  
data['author', 'affiliated_with', 'institution'].edge_attr = ...  
data['paper', 'has_topic', 'field_of_study'].edge_attr = ...  
这样上面的异质图就建立完成了，我们可以将它输入到一个异质图神经网络中

# 异质图神经网络
model = HeteroGNN(...)
# 获取异质图神经网络网络的输出
# 注意异质图神经网络的输入是 ..._dict
output = model(data.x_dict, data.edge_index_dict, data.edge_attr_dict)
如果PyG中包含你想用的异质图，可以直接这样导入

from torch_geometric.datasets import OGB_MAG

# 导入数据集
dataset = OGB_MAG(
    root='../data',
    # 预处理方式
    # 转换为向量
    preprocess='metapath2vec',
)
print(dataset[0])

# 异质图神经网络
model = HeteroGNN(...)
# 获取异质图神经网络网络的输出
# 注意异质图神经网络的输入是 ..._dict
output = model(data.x_dict, data.edge_index_dict, data.edge_attr_dict)

from torch_geometric.datasets import OGB_MAG

# 导入数据集
dataset = OGB_MAG(
    root='../data',
    # 预处理方式
    # 转换为向量
    preprocess='metapath2vec',
)
print(dataset[0])

#获取异质图中的某种节点或边
paper_node_data=data['paper']
cites_edge_data=data['paper','cites','paper']
#如果边的连接节点集合或者边的命名是唯一的还可以这样写
#使用连接端点获取边
cites_edge_data=data['paper','paper']
#使用边的名字获取
cites_edge_data=data['cites']
#给节点添加新属性
data['paper'].year=...
#删除节点的某些属性
def data['field_of_study']
#通过metadata获取异质图中所有类型的信息
node_types,edge_types=data.metadata()
#所有类型的节点
print(node_types)
#所有类型的边
print(edge_types)
#判断异质图自身的一些属性
print(data.has_isolated_nodes())
#如果不同类型信息之间维度匹配还可以将异质图融合为一个简单图
homogeneous_data=data.to_homogeneous()
import torch_geometric.transforms as T
#对异质图进行变换
#变为无向图
data=T.ToUndirected()(data)
#添加到自身的环
data=T.AddSelfLoops()(data)

import torch
import torch_geometric.transforms as T
from torch_geometric.datasets import OGB_MAG
from torch_geometric.nn import SAGEConv, to_hetero

#导入数据集
data = OGB_MAG(
    root='./data', 
    preprocess='metapath2vec', 
    transform=T.ToUndirected())[0]

#定义一个普通的图神经网络
class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        self.conv2 = SAGEConv((-1, -1), out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x

#实例化我们定义的图神经网络
model = GNN(hidden_channels=64, out_channels=dataset.num_classes)
#将其转换为异质图形式
model = to_hetero(model, data.metadata(), aggr='sum')

from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv, GATConv, Linear

class HeteroGNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, num_layers):
        super().__init__()

        self.convs = torch.nn.ModuleList()
        #定义图卷积层
        for _ in range(num_layers):
            #最外面用HeteroConv将里面的卷积层转换为异质图版本
            conv = HeteroConv(
                #要转换的卷积层
                {
                ('paper', 'cites', 'paper'): GCNConv(-1, hidden_channels),
                ('author', 'writes', 'paper'): GATConv((-1, -1), hidden_channels),
                ('author', 'affiliated_with', 'institution'): SAGEConv((-1, -1), hidden_channels),
                }, 
                aggr='sum')
            self.convs.append(conv)

        self.lin = Linear(hidden_channels, out_channels)

    def forward(self, x_dict, edge_index_dict):
        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict)
            x_dict = {key: x.relu() for key, x in x_dict.items()}
        return self.lin(x_dict['author'])

model = HeteroGNN(hidden_channels=64, out_channels=dataset.num_classes,
                  num_layers=2)

url='https://gitee.com/jiajiewu/planetoid/raw/master/data'

import osos.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'