PyTorch Geometric 图神经网络模型训练和特征文件加载

本文将介绍使用 PyTorch Geometric 库训练图神经网络模型，并演示如何加载新的特征文件，实现图像特征的提取和利用。

1. 数据加载和数据集创建

首先，我们定义一个 MyDataset 类来加载数据并创建 PyTorch Geometric 数据集。

import os
import pandas as pd
import torch
import torch.nn as nn
from torch_geometric.data import Data, DataLoader
from torch_geometric.nn import GCNConv
import torch.nn.functional as F
from sklearn.model_selection import train_test_split

# 加载数据并创建PyG数据集类：
class MyDataset(torch.utils.data.Dataset):
    def __init__(self, root, transform=None, pre_transform=None):
        self.edges = pd.read_csv(os.path.join(root, 'edges.csv'), header=None)
        self.features = pd.read_csv(os.path.join(root, 'features.csv'), header=None)
        self.labels = pd.read_csv(os.path.join(root, 'labels.csv'), header=None)
        self.transform = transform
        self.pre_transform = pre_transform

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        edge_index = torch.tensor(self.edges.values, dtype=torch.long).t().contiguous()
        features_file = f'C:\Users\jh\Desktop\data\images_flatten\{idx}.txt'
        with open(features_file, 'r') as f:
            features = [float(x) for x in f.read().split()]
        x = torch.tensor(features, dtype=torch.float).view(-1, 1)
        y = torch.tensor(self.labels.iloc[idx].values, dtype=torch.long)

        train_mask = torch.zeros(y.size(0), dtype=torch.bool)
        val_mask = torch.zeros(y.size(0), dtype=torch.bool)
        train_mask[:16] = 1
        val_mask[16:] = 1

        data = Data(x=x, edge_index=edge_index, y=y, train_mask=train_mask, val_mask=val_mask)

        if self.transform is not None:
            data = self.transform(data)

        return data

在修改后的 __getitem__ 方法中，通过 features_file 变量构建特征文件的路径，并使用 open 函数读取文件内容。然后将读取到的特征值转换为浮点数，并使用 view 方法将其 reshape 为一个列向量。最后将特征向量转换为 PyTorch 的张量。

2. 定义图神经网络模型

接下来，定义一个基于图卷积网络 (GCN) 的模型。

class GCN(torch.nn.Module):
    def __init__(self, num_node_features, num_classes):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(num_node_features, 32)
        self.conv2 = GCNConv(32, 64)
        self.conv3 = GCNConv(64, num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv3(x, edge_index)
        return x

3. 训练和验证模型

最后，定义训练和验证模型的函数，并加载数据集、创建模型、定义优化器和训练循环。

def train_model(dataset, model, optimizer, device):
    model.train()
    total_loss = 0.0

    for data in dataset:
        data = data.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output[data.train_mask], data.y[data.train_mask])
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    return total_loss / len(dataset)

def validate_model(dataset, model, device):
    model.eval()
    correct = 0
    total = 0

    for data in dataset:
        data = data.to(device)
        output = model(data)
        _, predicted = torch.max(output[data.val_mask], 1)
        total += data.val_mask.sum().item()
        correct += (predicted == data.y[data.val_mask]).sum().item()

    return correct / total

if __name__ == '__main__':
    dataset = MyDataset(root="C:\Users\jh\Desktop\data\raw1")
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    model = GCN(num_node_features=2, num_classes=3).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.002)

    train_dataset, val_dataset = train_test_split(dataset, test_size=0.2)
    train_loader = DataLoader(train_dataset, batch_size=1, shuffle=False)
    val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False)

    epochs = 2000
    for epoch in range(epochs):
        train_loss = train_model(train_loader, model, optimizer, device)
        print(f'Epoch {epoch+1}/{epochs}, Train Loss: {train_loss:.4f}')
        val_accuracy = validate_model(val_loader, model, device)
        print(f'Val_Acc: {val_accuracy:.4f}')

4. 总结

本文演示了如何使用 PyTorch Geometric 库训练图神经网络模型，并展示了如何加载新的特征文件，实现图像特征的提取和利用。这个示例可以帮助您更好地理解图神经网络的应用，并为您的图像识别和分析任务提供一些启发。

注意:

上述代码假设新的特征文件命名规则为i.txt_j.txt，其中i表示图片编号，j表示节点编号。你可以根据实际情况进行修改。
代码中对特征值文件的读取方式做了简化，你可以根据实际文件格式进行相应的修改。
代码中的参数值（如 epochs、lr 等）可以根据具体情况进行调整。

希望本文对您有所帮助。如果您有任何问题或建议，请随时在评论区留言。