pytorch mm 内存占用 pytorch需要多大内存

目录

1.整体介绍

1.1 图的表示

1.2 常用的图神经网络数据集

1.3 如何加载数据集

2.空域图卷积神经网络的建立

2.1 GCN的实现

2.2 Edge Convolution的实现

3.自建图神经网络数据集

3.1 创建一个能够完全放入内存中的图数据集

3.2 创建无法完全放入内存的数据集

4.批处理

5.异质图的建立

6.异质图神经网络的建立

6.1 将简单图神经网络转换为异质图神经网

---

---

PyG(PyTorch Geometric)是一个基于PyTorch的图神经网络框架，建议先了解PyTorch的使用再学习PyG

PyG支持图神经网络训练中的数据集处理、多GPU训练、多个经典的图神经网络模型、多个常用的图神经网络训练数据集而且支持自建数据集，主要包含以下几个模块

• torch_geometric:主模块
• torch_geometric.nn：搭建图神经网络层
• torch_geometric.data:图结构数据的表示
• torch_geometric.loader:加载数据集
• torch_geometric.datasets:常用的图神经网络数据集
• torch_geometric.transforms:数据变换
• torch_geometric.utils:常用工具
• torch_geometric.graphgym:常用的图神经网络模型
• torch_geometric.profile:监督模型的训练

1.整体介绍

1.1 图的表示

PyG用torch_geometric.data.Data保存图结构的数据，导入的data(这个data指的是你导入的具体数据，不是前面那个torch_geometric.data)在PyG中会包含以下属性

• data.x：图节点的属性信息，比如社交网络中每个用户是一个节点，这个x可以表示用户的属性信息，维度为[num_nodes,num_node_features]
• data.edge_index:COO格式的图节点连接信息，类型为torch.long,维度为[2,num_edges](具体包含两个列表，每个列表对应位置上的数字表示相应节点之间存在边连接）
• data.edge_attr:图中边的属性信息，维度[num_edges,num_edge_features]
• data.y:标签信息，根据具体任务，维度是不一样的，如果是在节点上的分类任务，维度为[num_edges,类别数],如果是在整个图上的分类任务，维度为[1,类别数]
• data.pos:节点的位置信息（一般用于图结构数据的可视化）

除了以上属性，我们还可以通过data.face自定义属性。

下面看如何利用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())

1.2 常用的图神经网络数据集

PyG包含了一些常用的图深度学习公共数据集，如

• Planetoid数据集（Cora、Citeseer、Pubmed）
• 一些来自于http://graphkernels.cs.tu-dortmund.de常用的图神经网络分类数据集
• QM7、QM9
• 3D点云数据集，如FAUST、ModelNet10等

接下来拿ENZYMES数据集（包含600个图，每个图分为6个类别，图级别的分类）举例如何使用PyG的公共数据集

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()

1.3 如何加载数据集

真正的图神经网络训练中我们一般是加载数据集中的一部分到内存中训练图神经网络，叫做一个batch,那么PyG如何加载一个batch呢，PyG会根据我们的数据集将其分割为我们指定的batch大小

举个例子

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)

2.空域图卷积神经网络的建立

空域图卷积（注意，图神经网络里的‘卷积’一词，取得是‘特征提取’这个广义意义，跟卷积神经网络里的那个卷积核计算区别开）可以看作是相邻节点之间进行信息传递、融合的过程，计算公式可以一般化为

 

对于以上计算过程,PyG利用MessagePassing进行实现。接下来以两篇经典图神经网络论文为例，介绍MessagePassing的使用

https://arxiv.org/abs/1609.02907

https://arxiv.org/abs/1801.07829

2.1 GCN的实现

在第一篇论文中，作者提出的卷积计算公式为

mport 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)

2.2 Edge Convolution的实现

在第二篇论文中，作者提出的卷积公式为

 

 

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)

3.自建图神经网络数据集

PyG将自建数据集分为两个文件夹---raw_dir、processed_dir。row_dir是原始的数据集，processed_dir是PyG处理之后的数据集

对于数据集PyG有三种过滤方法---transform、pre_transform、pre_filter。

• transform:读取数据，然后对其进行变换
• pre_transform：对于整个数据集进行变换，然后将变换之后的数据进行存储，pre_filter同理

PyG将数据集分为两种类型

• torch_geometric.data.InMemoryDataset：能够完全放入内存中的
• torch_geometric.data.Dataset:不能够完全放入内存中的

3.1 创建一个能够完全放入内存中的图数据集

做4件事：

1. 实现torch_geometric.data.InMemoryDataset.raw_file_names():告诉PyG数据集放在哪里
2. 实现torch_geometric.data.InMemoryDataset.processed_file_names():告诉PyG数据集处理完之后放在哪里
3. 实现torch_geometric.data.InMemoryDataset.download():告诉PyG从哪里获取数据集
4. 实现torch_geometric.data.InMemoryDataset.process():告诉PyG如何处理你的数据集

一个通用模板是这样的

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])

3.2 创建无法完全放入内存的数据集

这种就类似于PyTorch中的Dataset了，在上面需要做的几件事的基础上还需要

1. 实现torch_geometric.data.Dataset.len():告诉PyG数据集有多大
2. 实现torch_geometric.data.Dataset.get():告诉PyG如何从数据集中获取一个数据

通用模板为

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

4.批处理

来源于传统深度学习中批处理的思想---将数据分批，然后每批数据组合为一组，然后一组、一组的进行训练，每组数据的数据量叫做batch_size。PyG是将图数据集分为多组进行训练的

PyG会自动帮我们将图数据集按照我们定义的batch_size分割，然后将每个batch中的数据合并。

如果我们想要控制PyG如何组合一个batch中的数据，我们需要自己重写torch_geometric.data.Data.__inc__()

举两个具体的例子

假设我们的数据集中每个数据（注意是每个数据）包含两个图，每个数据像这样

对于这种数据集，如何控制PyG将多个数据合并成一个batch呢，以batch_size为2举例，batch_size=2意味着将数据集中每两个数据组成一组，形成一个图，每个batch中的数据是这样的

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)

再举一个二分图的例子，假设我们数据集中每个数据是一个二分图，像这样

 还是batch_size=2,我们想控制PyG让数据变成

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)

5.异质图的建立

前面讨论的图可以归为简单图---只包含一种类型的节点以及一种类型的边。

然而在现实中需要对多种类型的节点以及这些节点之间多种类型的边进行处理，这就需要异质图的概念，在异质图中，不同类型的边 描述不同类型节点之间 不同的关系，异质图神经网络的任务就是在这种图结构上学习出节点或者整个异质图的特征表示。异质图准确定义如下：

接下来以一个电影评分数据集MovieLens举例如何构建异质图。

MovieLens包含了600个用户对于电影的评分，我们利用这个数据集构建一个二分图，包含电影、用户两种类型的节点，一种类型的边（含有多种类型节点，所以可以看作一个异质图）

MovieLens中的movies.csv文件描述了电影的信息，包括电影在数据集中唯一的ID，电影名，电影所属的类型

 ratings.csv包含了用户对于电影的评分

 接下来就根据这两个csv建立二分图数据集

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)

6.异质图神经网络的建立

拿OGB数据集举例

在OGB数据集中包含4种类型的节点

1. author
2. paper
3. institution
4. field of study

4种类型的边

1. writes:author和paper之间的连接关系
2. affiliated with:author和institution之间的连接关系
3. cites:paper和paper之间的关系
4. has topic:paper和field of study之间的关系

OGB数据集上的任务是预测论文在整个关系网中所属的位置

下面看如何表示这个异质图

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])

下面介绍一下HeteroData中常用的函数

#获取异质图中的某种节点或边
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)

下面介绍如何建立异质图神经网络

6.1 将简单图神经网络转换为异质图神经网络

PyG可以通过torch_geometric.nn.to_hetero(),或者torch_geometric.nn.to_hetero_with_bases()将一个简单图神经网络转换成异质图的形式

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')

PyG的to_hetero具体工作方式是这样的

它根据我们的异质图数据结构，自动将我们定义的简单图神经网络结构中的层结构进行了复制，并添加了信息传递路径。

torch_geometric.nn.conv.HeteroConv卷积层同样起到类似的功能

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)