PyTorch PyG能支持自定义层吗

PyTorch的PyG库可以支持自定义层。在PyTorch中，可以通过继承torch.nn.Module类来创建自定义层。例如，定义一个简单的全连接层，可以这样做：

import torch
import torch.nn as nn

class MyLayer(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(MyLayer, self).__init__()
        self.linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        return self.linear(x)

在这个例子中，MyLayer类继承自nn.Module，并定义了一个全连接层self.linear。在forward方法中，我们将输入x传递给这个全连接层，并返回其输出。

然后，在使用PyG库时，可以将这个自定义层添加到图结构中。例如，定义一个包含自定义层和PyTorch nn.Linear层的图结构：

from torch_geometric.nn import MessagePassing
import torch

class MyModel(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super(MyModel, self).__init__(aggr='add')
        self.lin = nn.Linear(in_channels, out_channels)
        self.my_layer = MyLayer(in_channels, 64)

    def forward(self, x, edge_index):
        row, col = edge_index
        x = self.my_layer(x)
        x = self.lin(x)
        row, col = row.view(-1, 1), col.view(-1, 1)
        deg = self.degree(row, x.size(0), dtype=x.dtype)
        deg_inv_sqrt = deg.pow(-0.5)
        norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]
        return self.propagate(edge_index, x=x, norm=norm)

    def message(self, x_j, norm):
        return norm.view(-1, 1) * x_j

    def degree(self, row, num_nodes, dtype):
        row, col = row.to(dtype), col.to(dtype)
        deg = torch.bincount(row, minlength=num_nodes, dtype=dtype)
        deg = deg[row] + deg[col]
        return deg.view(-1, 1)

在这个例子中，MyModel类继承自MessagePassing，并定义了一个包含自定义层self.my_layer和PyTorch nn.Linear层的图结构。在forward方法中，我们首先对输入x应用自定义层，然后应用线性层，最后根据边的权重计算消息和更新节点特征。