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

"""import torch\nimport torch.nn.functional as F\nfrom torch_geometric.data import Data\nfrom torch_geometric.data import DataLoader\nfrom torch_geometric.nn import GCNConv\nfrom torch_geometric.datasets import InMemoryDataset\n\nclass MyDataset(InMemoryDataset):\n def init(self, root, transform=None, pre_transform=None):\n super(MyDataset, self).init(root, transform, pre_transform)\n self.data, self.slices = torch.load(self.processed_paths[0])\n \n @property\n def raw_file_names(self):\n return ['edges.csv', 'features1.csv', 'features2.csv', 'label.csv']\n \n @property\n def processed_file_names(self):\n return ['data.pt']\n \n def download(self):\n # Download the raw data to the self.raw_dir folder.\n pass\n \n def process(self):\n # Read raw data from files and process into PyG Data objects.\n edges_path = self.raw_paths[0]\n features1_path = self.raw_paths[1]\n features2_path = self.raw_paths[2]\n label_path = self.raw_paths[3]\n \n # Process edges\n edge_index = []\n with open(edges_path, 'r') as f:\n for line in f:\n src, dst = line.strip().split(',')\n edge_index.append([int(src), int(dst)])\n \n # Process features1\n features1 = []\n with open(features1_path, 'r') as f:\n for line in f:\n features1.append([float(x) for x in line.strip().split(',')])\n \n # Process features2\n features2 = []\n with open(features2_path, 'r') as f:\n for line in f:\n features2.append([float(x) for x in line.strip().split(',')])\n \n # Process labels\n labels = []\n with open(label_path, 'r') as f:\n for line in f:\n labels.append(int(line.strip()))\n \n x = torch.tensor(features1 + features2, dtype=torch.float)\n edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()\n y = torch.tensor(labels, dtype=torch.long)\n \n data = Data(x=x, edge_index=edge_index, y=y)\n data = data if self.pre_transform is None else self.pre_transform(data)\n torch.save(self.collate([data]), self.processed_paths[0])\n\nclass GCN(torch.nn.Module):\n def init(self):\n super(GCN, self).init()\n self.conv1 = GCNConv(2, 16)\n self.conv2 = GCNConv(16, 2)\n \n def forward(self, data):\n x, edge_index = data.x, data.edge_index\n x = self.conv1(x, edge_index)\n x = F.relu(x)\n x = self.conv2(x, edge_index)\n return F.log_softmax(x, dim=1)\n\n# Create dataset\ndataset = MyDataset(root='C:\Users\jh\Desktop\data')\n\n# Create model\nmodel = GCN()\n\n# Create train and validation masks\ntrain_mask = torch.zeros((10, 20), dtype=torch.uint8)\ntrain_mask[:, :16] = 1\nval_mask = 1 - train_mask\n\n# Create data loaders\ntrain_loader = DataLoader(dataset[train_mask], batch_size=1, shuffle=True)\nval_loader = DataLoader(dataset[val_mask], batch_size=1, shuffle=False)\n\n# Define optimizer and loss function\noptimizer = torch.optim.Adam(model.parameters(), lr=0.01)\ncriterion = torch.nn.NLLLoss()\n\n# Training loop\nfor epoch in range(10):\n model.train()\n for batch in train_loader:\n optimizer.zero_grad()\n out = model(batch)\n loss = criterion(out, batch.y)\n loss.backward()\n optimizer.step()\n \n model.eval()\n correct = 0\n total = 0\n for batch in val_loader:\n out = model(batch)\n _, predicted = torch.max(out, 1)\n total += batch.y.size(0)\n correct += (predicted == batch.y).sum().item()\n \n accuracy = 100 * correct / total\n print(f'Epoch: {epoch+1}, Accuracy: {accuracy}%')\n"""