Pytorch+PyG实现GraphConv过程示例详解

# 导入cifar-10数据集
from torch_geometric.datasets import Planetoid
# 加载数据、划分训练集和测试集
dataset = Planetoid(root='./cifar10', name='Cora')
data = dataset[0]
# 定义超级参数
num_features = dataset.num_features
num_classes = dataset.num_classes
# 构建训练集和测试集索引文件
train_mask = torch.zeros(data.num_nodes, dtype=torch.uint8)
train_mask[:800] = 1
test_mask = torch.zeros(data.num_nodes, dtype=torch.uint8)
test_mask[800:] = 1
# 创建数据加载器
train_loader = DataLoader(data[train_mask], batch_size=32, shuffle=True)
test_loader = DataLoader(data[test_mask], batch_size=32, shuffle=False)

from torch.nn import Linear, ModuleList, ReLU
from torch_geometric.nn import GCNConv
class GraphConv(torch.nn.Module):
    def __init__(self, dataset):
        super(GraphConv, self).__init__()
        # 定义基础参数
        self.input_dim = dataset.num_features
        self.output_dim = dataset.num_classes
        # 定义GCN网络结构
        self.convs = ModuleList()
        self.convs.append(GCNConv(self.input_dim, 16))
        self.convs.append(GCNConv(16, 32))
        self.convs.append(GCNConv(32, self.output_dim))
    def forward(self, x, edge_index):
        for conv in self.convs:
            x = conv(x, edge_index)
            x = F.relu(x)
        return F.log_softmax(x, dim=1)

# 定义训练计划，包括损失函数、优化器及迭代次数等
train_epochs = 200
learning_rate = 0.01
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(graph_conv.parameters(), lr=learning_rate)
losses_per_epoch = []
accuracies_per_epoch = []
for epoch in range(train_epochs):
    running_loss = 0.0
    running_corrects = 0.0
    count = 0.0
    for samples in train_loader:
        optimizer.zero_grad()
        x, edge_index = samples.x, samples.edge_index
        out = graph_conv(x, edge_index)
        label = samples.y
        loss = criterion(out, label)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() / len(train_loader.dataset)
        pred = out.argmax(dim=1)
        running_corrects += pred.eq(label).sum().item() / len(train_loader.dataset)
        count += 1
    losses_per_epoch.append(running_loss)
    accuracies_per_epoch.append(running_corrects)
    if (epoch + 1) % 20 == 0:
        print("Train Epoch {}/{} Loss {:.4f} Accuracy {:.4f}".format(
            epoch + 1, train_epochs, running_loss, running_corrects))