PyG自定义数据集：基于节点连接关系和特征的图神经网络模型训练与验证

以下是使用 PYG 库自定义数据集 My_dataset 类的示例代码，该代码演示了如何利用节点连接关系和特征信息来训练图神经网络模型，并计算训练过程中的损失和准确率。

import torch
import torch_geometric.transforms as T
from torch_geometric.data import DataLoader
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv
import torch.nn.functional as F

from sklearn.metrics import accuracy_score

from torch_geometric.data import Data

class My_dataset(Data):
    def __init__(self, root, transform=None, pre_transform=None):
        super(My_dataset, self).__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return []

    @property
    def processed_file_names(self):
        return ['data.pt']

    def download(self):
        pass

    def process(self):
        edges = torch.tensor([
            [0, 1, 0, 2, 0, 3, 0, 4, 1, 2, 1, 3, 1, 9, 2, 3, 2, 5, 2, 8, 3, 6, 3, 8, 3, 10, 3, 13, 3, 15, 3, 18, 4, 5, 4, 9, 4, 10, 5, 11, 5, 12, 5, 13, 6, 16, 6, 17, 6, 19, 7, 8, 7, 13, 8, 13, 8, 18, 9, 10, 9, 16, 10, 11, 10, 15, 10, 16, 11, 13, 12, 14, 12, 16, 13, 19, 14, 15, 14, 16, 17, 18, 17, 19, 18, 19, 5, 18, 6, 8, 14, 18, 12, 15, 12, 18, 13, 14, 11, 12, 11, 14, 11, 19, 8, 9, 8, 11, 8, 15, 7, 19, 7, 9, 6, 10, 6, 13, 7, 10, 13, 17, 7, 14, 7, 17, 9, 12, 9, 18, 5, 16]
        ], dtype=torch.long).view(2, -1)

        x1 = torch.tensor([
            [1, 0, 1, 1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0],
            [1, 0, 1, 1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0],
        ], dtype=torch.float)

        x2 = torch.tensor([
            [1, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0],
            [1, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0],
        ], dtype=torch.float)

        y = torch.tensor([
            [0, 0, 1, 1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0],
            [0, 0, 1, 1, 1, 1, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0],
        ], dtype=torch.float)

        data = Data(x1=x1, x2=x2, y=y, edge_index=edges)
        data = data if self.pre_transform is None else self.pre_transform(data)
        torch.save(self.collate([data]), self.processed_paths[0])

dataset = My_dataset(root='.', transform=T.NormalizeFeatures())
loader = DataLoader(dataset, batch_size=1, shuffle=True)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = GCNConv(dataset.num_node_features, 16)
        self.conv2 = GCNConv(16, dataset.num_node_features)

    def forward(self, data):
        x1, x2 = data.x1, data.x2
        edge_index = data.edge_index

        x1 = self.conv1(x1, edge_index)
        x1 = F.relu(x1)
        x1 = self.conv2(x1, edge_index)
        x1 = F.relu(x1)

        x2 = self.conv1(x2, edge_index)
        x2 = F.relu(x2)
        x2 = self.conv2(x2, edge_index)
        x2 = F.relu(x2)

        return x1, x2

model = Net().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

def train():
    model.train()
    total_loss = 0
    for data in loader:
        data = data.to(device)
        optimizer.zero_grad()
        out1, out2 = model(data)
        loss = F.mse_loss(out1, data.y) + F.mse_loss(out2, data.y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * data.num_graphs
    return total_loss / len(dataset)

def test():
    model.eval()
    y_true = []
    y_pred = []
    for data in loader:
        data = data.to(device)
        with torch.no_grad():
            out1, out2 = model(data)
            pred = torch.cat([out1, out2], dim=0)
            y_true.append(data.y.view(-1).cpu().numpy())
            y_pred.append(pred.view(-1).cpu().numpy())
    y_true = np.concatenate(y_true)
    y_pred = np.concatenate(y_pred)
    acc = accuracy_score(y_true, np.round(y_pred))
    return acc

for epoch in range(1, 201):
    loss = train()
    acc = test()
    print(f'Epoch: {epoch}, Loss: {loss:.4f}, Accuracy: {acc:.4f}')

该代码定义了一个名为 My_dataset 的自定义数据集类，继承自 Data 类。在 process 方法中，根据给定的边和特征信息构造了一个 Data 对象，并将其保存为 .pt 文件。在训练过程中，利用 GCNConv 层对节点特征进行卷积操作，然后使用均方误差损失函数计算损失，并根据损失进行反向传播和参数更新。在验证过程中，将模型的输出结果与真实标签进行对比，计算准确率。最后，将训练过程中的损失和准确率输出。

代码示例中包含以下关键内容：

• 自定义数据集类 My_dataset
• 数据加载器 DataLoader
• 图卷积神经网络模型 Net
• 训练函数 train
• 验证函数 test
• 训练循环

该代码示例可以帮助你理解如何在 PyG 库中创建自定义数据集，并使用图神经网络模型进行训练和验证。你可以根据自己的需要修改代码，例如更改模型结构、添加新的特征信息等。