使用PyTorch写一个图神经网络

这里提供一个简单的图神经网络的PyTorch实现，用于节点分类任务。

首先，我们需要导入必要的库：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
import torch_geometric.nn as pyg_nn
import torch_geometric.transforms as T

然后，我们需要加载数据集。这里使用了Planetoid数据集中的Cora数据集。

dataset = Planetoid(root='/data', name='Cora', transform=T.NormalizeFeatures())
data = dataset[0]

接下来，我们定义一个简单的图神经网络模型。这个模型包含一个GCN层和一个线性层。

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

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

在forward函数中，我们首先将节点特征x和边信息edge_index传入GCN层中，得到GCN层的输出。然后，我们对输出进行ReLU激活和dropout操作，最后将其输入线性层进行分类。

接下来，我们定义训练函数和测试函数。

def train(model, optimizer, data):
    model.train()
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()

def test(model, data):
    model.eval()
    out = model(data)
    pred = out.argmax(dim=1)
    acc = pred[data.test_mask].eq(data.y[data.test_mask]).sum().item() / data.test_mask.sum().item()
    return acc

在训练函数中，我们首先将模型设为训练模式，然后使用优化器将梯度清零。接着，我们将数据传入模型中，得到输出并计算损失。最后，我们反向传播并更新模型参数。

在测试函数中，我们首先将模型设为测试模式，然后将数据传入模型中得到输出。接着，我们将输出转换为预测结果，并计算准确率。

最后，我们开始训练和测试模型。

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

for epoch in range(200):
    loss = train(model, optimizer, data)
    acc = test(model, data)
    print('Epoch {}, Loss: {:.4f}, Test Acc: {:.4f}'.format(epoch, loss, acc))

这里我们使用Adam优化器和交叉熵损失函数，进行200个epoch的训练。每个epoch结束后，我们输出当前的损失和测试准确率。

完整代码如下：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
import torch_geometric.nn as pyg_nn
import torch_geometric.transforms as T

dataset = Planetoid(root='/data', name='Cora', transform=T.NormalizeFeatures())
data = dataset[0]

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

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

def train(model, optimizer, data):
    model.train()
    optimizer.zero_grad()
    out = model(data)
    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()

def test(model, data):
    model.eval()
    out = model(data)
    pred = out.argmax(dim=1)
    acc = pred[data.test_mask].eq(data.y[data.test_mask]).sum().item() / data.test_mask.sum().item()
    return acc

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

for epoch in range(200):
    loss = train(model, optimizer, data)
    acc = test(model, data)
    print('Epoch {}, Loss: {:.4f}, Test Acc: {:.4f}'.format(epoch, loss, acc))