我用networkx构蛋白质图，每个残基作为一个节点，并给每个节点添加了一个19维德特征，讲图保存为了pkl文件，现在我想用pyg将图导入，用于下一步图神经网络的训练，请给出导入数据的代码以及模型搭建、训练的示例代码

在PyTorch Geometric (PyG)库中，你可以使用`torch_geometric.data.Data`类来加载你的蛋白图数据，包括节点特征。首先，你需要安装必要的库：

pip install torch torchvision networkx torch-scatter -f https://data.pyg.org/whl/torch-1.8.0+cu111.html

假设你已经有一个名为`protein_graph.pkl`的pkl文件，其中包含`edge_index`, `node_features`, 和其他可能的元数据（如节点标签）。以下是加载数据并创建模型进行训练的示例代码：

import torch
import torch.nn as nn
from torch_geometric.datasets import ProteinGraphDataset
from torch_geometric.data import Data
from torch_geometric.utils import from_networkx

# 加载数据
class ProteinDataset(torch.utils.data.Dataset):
    def __init__(self, path):
        self.data = Data.from_dict(torch.load(path))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

dataset = ProteinDataset('protein_graph.pkl')
train_dataset, test_dataset = torch.utils.data.random_split(dataset, [int(len(dataset) * 0.8), len(dataset) * 0.2])

# 创建模型（这里是一个简单的图卷积网络）
class ProteinGraphModel(nn.Module):
    def __init__(self):
        super(ProteinGraphModel, self).__init__()
        self.conv1 = GCNConv(dataset.num_node_features, 64)
        self.conv2 = GCNConv(64, dataset.num_classes)  # 假设你的目标分类有dataset.num_classes类别

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

model = ProteinGraphModel()

# 定义损失函数和优化器
criterion = nn.NLLLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 训练循环
for epoch in range(10):  # 更换为实际需要的迭代次数
    for data in train_dataset:
        optimizer.zero_grad()
        out = model(data)
        loss = criterion(out[data.train_mask], data.y[data.train_mask])
        loss.backward()
        optimizer.step()

    # 评估阶段
    with torch.no_grad():
        z = model(test_dataset[0]).max(1)[1]
        print(f"Epoch {epoch + 1}: Train Loss: {loss.item():.4f}, Test Acc: {z.eq(test_dataset[0].y[test_dataset.test_mask]).sum().item() / test_dataset[0].y[test_dataset.test_mask].size(0):.4f}")

#