为了实现这个需求,首先需要安装 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}')

需要注意的是,上述代码中的一些细节(如数据路径、优化器的选择等)可能需要根据具体情况进行调整。

基于Torch Geometric 的节点分类模型训练与验证

原文地址: https://www.cveoy.top/t/topic/mt6O 著作权归作者所有。请勿转载和采集!

免费AI点我,无需注册和登录