基于Torch Geometric 的节点分类模型训练与验证
为了实现这个需求,首先需要安装 torch_geometric 库。可以使用以下命令进行安装:
pip install torch-geometric
接下来,我们可以定义自定义数据集 My_dataset 类,并根据要求加载数据。
import torch
from torch_geometric.data import Data, Dataset
class My_dataset(Dataset):
def __init__(self, root):
self.root = root
# 加载边数据
edges_path = self.processed_paths[0]
self.edges = torch.tensor(self.read_csv(edges_path), dtype=torch.long).t()
# 加载特征数据
features_path = [self.processed_paths[1], self.processed_paths[2]]
self.features = [torch.tensor(self.read_csv(path), dtype=torch.float) for path in features_path]
# 加载标签数据
label_path = self.processed_paths[3]
self.labels = torch.tensor(self.read_csv(label_path), dtype=torch.long)
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
edge_index = self.edges
x = self.features[:, idx, :]
y = self.labels[idx]
# 根据要求将前 16 个节点作为训练集,剩下 4 个节点作为验证集
train_mask = torch.zeros(x.shape[0], dtype=torch.bool)
train_mask[:16] = 1
# 创建 Data 对象用于训练和验证
data = Data(x=x, edge_index=edge_index, y=y, train_mask=train_mask)
return data
@property
def raw_file_names(self):
return ['edges.csv', 'features1.csv', 'features2.csv', 'label.csv']
@property
def processed_file_names(self):
return ['processed_edges.csv', 'processed_features1.csv', 'processed_features2.csv', 'processed_label.csv']
def download(self):
# 下载数据集
pass
def process(self):
# 数据预处理,将原始数据转换为 pytorch 所需的格式
for raw_path, processed_path in zip(self.raw_paths, self.processed_paths):
data = self.read_csv(raw_path)
self.save_csv(data, processed_path)
接下来,我们可以定义 GCN 模型,并进行训练和验证。
import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
class GCN(torch.nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
super(GCN, self).__init__()
self.convs = torch.nn.ModuleList()
self.convs.append(GCNConv(input_dim, hidden_dim))
for _ in range(num_layers - 2):
self.convs.append(GCNConv(hidden_dim, hidden_dim))
self.convs.append(GCNConv(hidden_dim, output_dim))
def forward(self, data):
x, edge_index = data.x, data.edge_index
for conv in self.convs:
x = F.relu(conv(x, edge_index))
return x
# 创建模型
model = GCN(input_dim=2, hidden_dim=16, output_dim=2, num_layers=3)
# 模型训练和验证
for epoch in range(num_epochs):
for data in dataset:
optimizer.zero_grad()
output = model(data)
loss = F.cross_entropy(output[data.train_mask], data.y[data.train_mask])
loss.backward()
optimizer.step()
with torch.no_grad():
val_loss = F.cross_entropy(output[~data.train_mask], data.y[~data.train_mask])
val_acc = (output[~data.train_mask].argmax(dim=1) == data.y[~data.train_mask]).sum().item() / (~data.train_mask).sum().item()
print(f'Epoch: {epoch}, Loss: {loss:.4f}, Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}')
需要注意的是,上述代码中的一些细节(如数据路径、优化器的选择等)可能需要根据具体情况进行调整。
原文地址: https://www.cveoy.top/t/topic/mt6O 著作权归作者所有。请勿转载和采集!