TorchScript 支持

TorchScript 是一种从 PyTorch 代码创建可序列化和可优化模型的方法。任何 TorchScript 程序都可以从 Python 进程中保存，并在没有 Python 依赖的进程中加载。如果您不熟悉 TorchScript，我们建议先阅读官方的“TorchScript 简介”教程。

转换GNN模型

注意

从 PyG 2.5（及以后版本）开始，GNN 层现在完全兼容 torch.jit.script()，无需任何修改。如果您使用的是较早版本的 PyG，请考虑通过调用 jittable() 将您的 GNN 层转换为“可 jit 化”实例。

将您的PyG模型转换为TorchScript程序非常简单，只需要进行少量代码更改。让我们考虑以下模型：

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

class GNN(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, 64)
        self.conv2 = GCNConv(64, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        return F.log_softmax(x, dim=1)

model = GNN(dataset.num_features, dataset.num_classes)

实例化的模型现在可以直接传递给 torch.jit.script()：

model = torch.jit.script(model)

这就是你需要知道的关于如何将你的PyG模型转换为TorchScript程序的所有内容。你可以进一步查看我们的JIT示例，这些示例展示了如何为节点和图分类模型获取TorchScript程序。

创建可即时编译的GNN操作符

所有 PyG MessagePassing 操作符都经过测试，可以转换为 TorchScript 程序。然而，如果你希望自己的 GNN 模块与 torch.jit.script() 兼容，你需要考虑以下两点：

正如人们所期望的那样，您的forward()代码可能需要进行调整，以便它满足TorchScript编译器的要求，例如，通过添加类型注释。

你需要告诉MessagePassing模块你传递给它的propagate()函数的类型。这可以通过两种不同的方式实现：

在名为propagate_type的字典中声明传播参数的类型：

from typing import Optional
from torch import Tensor
from torch_geometric.nn import MessagePassing

class MyConv(MessagePassing):
    propagate_type = {'x': Tensor, 'edge_weight': Optional[Tensor] }

    def forward(
        self,
        x: Tensor,
        edge_index: Tensor,
        edge_weight: Optional[Tensor] = None,
    ) -> Tensor:
        return self.propagate(edge_index, x=x, edge_weight=edge_weight)

在你的模块中将传播参数的类型声明为注释：

from typing import Optional
from torch import Tensor
from torch_geometric.nn import MessagePassing

class MyConv(MessagePassing):
    def forward(
        self,
        x: Tensor,
        edge_index: Tensor,
        edge_weight: Optional[Tensor] = None,
    ) -> Tensor:
        # propagate_type: (x: Tensor, edge_weight: Optional[Tensor])
        return self.propagate(edge_index, x=x, edge_weight=edge_weight)

如果没有给出这些选项，MessagePassing 模块将推断 propagate() 的参数类型为 torch.Tensor（模仿 TorchScript 为非注释参数推断的默认类型）。