PyG 图神经网络：自定义数据集和模型训练与验证

from\u0020torch_geometric.data\u0020import\u0020Dataset import\u0020torch

class\u0020My_dataset(Dataset): \u0020\u0020def\u0020__init__(self): \u0020\u0020\u0020\u0020super(My_dataset, self).init() \u0020\u0020\u0020\u0020#\u0020读取边数据 \u0020\u0020\u0020\u0020with\u0020open('edges.txt', 'r')\u0020as\u0020f: \u0020\u0020\u0020\u0020\u0020\u0020edges\u0020=\u0020f.readlines()[0].split() \u0020\u0020\u0020\u0020edges\u0020=\u0020list(map(int, edges)) \u0020\u0020\u0020\u0020#\u0020读取特征数据 \u0020\u0020\u0020\u0020with\u0020open('features1.txt', 'r')\u0020as\u0020f: \u0020\u0020\u0020\u0020\u0020\u0020features1\u0020=\u0020f.readlines() \u0020\u0020\u0020\u0020features1\u0020=\u0020[list(map(int, line.split()))\u0020for\u0020line\u0020in\u0020features1] \u0020\u0020\u0020\u0020with\u0020open('features2.txt', 'r')\u0020as\u0020f: \u0020\u0020\u0020\u0020\u0020\u0020features2\u0020=\u0020f.readlines() \u0020\u0020\u0020\u0020features2\u0020=\u0020[list(map(int, line.split()))\u0020for\u0020line\u0020in\u0020features2] \u0020\u0020\u0020\u0020#\u0020读取标签数据 \u0020\u0020\u0020\u0020with\u0020open('label.txt', 'r')\u0020as\u0020f: \u0020\u0020\u0020\u0020\u0020\u0020labels\u0020=\u0020f.readlines() \u0020\u0020\u0020\u0020labels\u0020=\u0020[list(map(int, line.split()))\u0020for\u0020line\u0020in\u0020labels] \u0020\u0020\u0020\u0020#\u0020创建图数据 \u0020\u0020\u0020\u0020self.graphs\u0020=\u0020[] \u0020\u0020\u0020\u0020for\u0020i\u0020in\u0020range(10): \u0020\u0020\u0020\u0020\u0020\u0020edge_index\u0020=\u0020torch.tensor(edges[i*40:(i+1)*40], dtype=torch.long).reshape(2, -1) \u0020\u0020\u0020\u0020\u0020\u0020x1\u0020=\u0020torch.tensor(features1[i], dtype=torch.float).reshape(-1, 2) \u0020\u0020\u0020\u0020\u0020\u0020x2\u0020=\u0020torch.tensor(features2[i], dtype=torch.float).reshape(-1, 2) \u0020\u0020\u0020\u0020\u0020\u0020y\u0020=\u0020torch.tensor(labels[i], dtype=torch.long).reshape(-1, 1) \u0020\u0020\u0020\u0020\u0020\u0020self.graphs.append((edge_index, x1, x2, y)) \u0020\u0020 \u0020\u0020def\u0020len(self): \u0020\u0020\u0020\u0020return\u0020len(self.graphs) \u0020\u0020 \u0020\u0020def\u0020get(self, idx): \u0020\u0020\u0020\u0020return\u0020self.graphs[idx]

import\u0020torch import\u0020torch.nn.functional\u0020as\u0020F from\u0020torch_geometric.nn\u0020import\u0020GCNConv from\u0020torch_geometric.data\u0020import\u0020DataLoader from\u0020sklearn.metrics\u0020import\u0020accuracy_score

#\u0020定义模型 class\u0020MyModel(torch.nn.Module): \u0020\u0020def\u0020__init__(self): \u0020\u0020\u0020\u0020super(MyModel, self).init() \u0020\u0020\u0020\u0020self.conv1\u0020=\u0020GCNConv(2, 16) \u0020\u0020\u0020\u0020self.conv2\u0020=\u0020GCNConv(16, 32) \u0020\u0020\u0020\u0020self.fc\u0020=\u0020torch.nn.Linear(32, 2) \u0020\u0020 \u0020\u0020def\u0020forward(self, x1, x2, edge_index): \u0020\u0020\u0020\u0020x1\u0020=\u0020self.conv1(x1, edge_index) \u0020\u0020\u0020\u0020x1\u0020=\u0020F.relu(x1) \u0020\u0020\u0020\u0020x1\u0020=\u0020F.dropout(x1, training=self.training) \u0020\u0020\u0020\u0020x1\u0020=\u0020self.conv2(x1, edge_index) \u0020\u0020\u0020\u0020 \u0020\u0020\u0020\u0020x2\u0020=\u0020self.conv1(x2, edge_index) \u0020\u0020\u0020\u0020x2\u0020=\u0020F.relu(x2) \u0020\u0020\u0020\u0020x2\u0020=\u0020F.dropout(x2, training=self.training) \u0020\u0020\u0020\u0020x2\u0020=\u0020self.conv2(x2, edge_index) \u0020\u0020\u0020\u0020 \u0020\u0020\u0020\u0020x\u0020=\u0020torch.cat((x1, x2), dim=1) \u0020\u0020\u0020\u0020x\u0020=\u0020F.relu(x) \u0020\u0020\u0020\u0020x\u0020=\u0020F.dropout(x, training=self.training) \u0020\u0020\u0020\u0020x\u0020=\u0020self.fc(x) \u0020\u0020\u0020\u0020 \u0020\u0020\u0020\u0020return\u0020F.log_softmax(x, dim=1)

#\u0020定义训练函数 def\u0020train(model, loader, optimizer, device): \u0020\u0020model.train() \u0020\u0020loss_all\u0020=\u00200 \u0020\u0020for\u0020data\u0020in\u0020loader: \u0020\u0020\u0020\u0020data\u0020=\u0020data.to(device) \u0020\u0020\u0020\u0020optimizer.zero_grad() \u0020\u0020\u0020\u0020output\u0020=\u0020model(data.x1, data.x2, data.edge_index) \u0020\u0020\u0020\u0020loss\u0020=\u0020F.nll_loss(output[data.train_mask], data.y[data.train_mask]) \u0020\u0020\u0020\u0020loss.backward() \u0020\u0020\u0020\u0020optimizer.step() \u0020\u0020\u0020\u0020loss_all\u0020+=\u0020loss.item() * data.num_nodes \u0020\u0020return\u0020loss_all / len(loader.dataset)

#\u0020定义验证函数 def\u0020test(model, loader, device): \u0020\u0020model.eval() \u0020\u0020correct\u0020=\u00200 \u0020\u0020for\u0020data\u0020in\u0020loader: \u0020\u0020\u0020\u0020data\u0020=\u0020data.to(device) \u0020\u0020\u0020\u0020output\u0020=\u0020model(data.x1, data.x2, data.edge_index) \u0020\u0020\u0020\u0020pred\u0020=\u0020output.argmax(dim=1) \u0020\u0020\u0020\u0020correct\u0020+=\u0020pred[data.val_mask].eq(data.y[data.val_mask].view_as(pred[data.val_mask])).sum().item() \u0020\u0020return\u0020correct / len(loader.dataset)

#\u0020设置设备 device\u0020=\u0020torch.device('cuda' if torch.cuda.is_available() else 'cpu')

#\u0020创建数据集对象 dataset\u0020=\u0020My_dataset()

#\u0020创建训练集和验证集 mask\u0020=\u0020torch.zeros(dataset[0][0].shape[1], dtype=torch.bool) mask[:16]\u0020=\u00201 train_dataset\u0020=\u0020dataset[mask] val_dataset\u0020=\u0020dataset[~mask]

#\u0020创建数据加载器 train_loader\u0020=\u0020DataLoader(train_dataset, batch_size=1, shuffle=True) val_loader\u0020=\u0020DataLoader(val_dataset, batch_size=1, shuffle=False)

#\u0020创建模型和优化器 model\u0020=\u0020MyModel().to(device) optimizer\u0020=\u0020torch.optim.Adam(model.parameters(), lr=0.01)

#\u0020训练和验证模型 for\u0020epoch\u0020in\u0020range(200): \u0020\u0020train_loss\u0020=\u0020train(model, train_loader, optimizer, device) \u0020\u0020val_acc\u0020=\u0020test(model, val_loader, device) \u0020\u0020print('Epoch: {:03d}, Train Loss: {:.4f}, Val Acc: {:.4f}'.format(epoch, train_loss, val_acc))