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.root = root
        self.transform = transform
        self.pre_transform = pre_transform

    def __len__(self):
        return len(os.listdir(os.path.join(self.root, 'images')))

    def __getitem__(self, idx):
        image_path = os.path.join(self.root, 'images', f'{idx + 1}.png')
        feature_path = os.path.join(self.root, 'features', f'{idx + 1}.txt')
        label_path = os.path.join(self.root, 'labels', f'{idx + 1}.txt')
        edge_path = os.path.join(self.root, 'edges_L.csv')

        image = self.load_image(image_path)
        feature = self.load_feature(feature_path)
        label = self.load_label(label_path)
        edge_index = self.load_edge(edge_path)

        data = Data(x=feature, edge_index=edge_index, y=label)

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

        return data

    def load_image(self, image_path):
        # Load and process image data
        # You can use any image processing library (e.g. PIL, OpenCV) here
        pass

    def load_feature(self, feature_path):
        # Load and process feature data
        # You can use any file reading library (e.g. pandas) here
        pass

    def load_label(self, label_path):
        # Load and process label data
        # You can use any file reading library (e.g. pandas) here
        pass

    def load_edge(self, edge_path):
        # Load and process edge data
        # You can use any file reading library (e.g. pandas) here
        pass

# 定义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

# 创建训练和验证模型:
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.y)
        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, 1)
        total += data.y.size(0)
        correct += (predicted == data.y).sum().item()

    return correct / total

# 加载数据集、创建模型、定义优化器和训练循环,以及验证模型:
if __name__ == '__main__':
    dataset = MyDataset(root="C:\Users\jh\Desktop\data\input")
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    model = GCN(num_node_features=8, num_classes=8).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}')

本项目使用图神经网络(GCN)对42个时刻的图数据进行分类,每个时刻包含37个节点,节点特征为图像的颜色特征,标签为8个类别,数据存储在指定的文件夹中。

数据存储结构:

  • images: 包含所有时刻的图像,文件名格式为 i.png,其中 i 表示时刻序号(从1到42)。
  • features: 包含每个时刻所有节点的颜色特征,文件名格式为 i.txt,其中 i 表示时刻序号。每个文件包含37行,每行代表一个节点的8个颜色特征值,用空格隔开。
  • labels: 包含每个时刻所有节点的标签,文件名格式为 i.txt,其中 i 表示时刻序号。每个文件包含37行,每行代表一个节点的8个标签值,用空格隔开。
  • edges_L.csv: 包含所有时刻的边连接关系,第一列为源节点,第二列为目标节点,边为无向边。

代码说明:

  • MyDataset 类用于加载和处理数据,并将其转换为 PyG 数据格式。
  • GCN 类定义了图卷积神经网络模型。
  • train_model 函数用于训练模型。
  • validate_model 函数用于评估模型性能。

运行步骤:

  1. 将数据按照上述结构存储在指定文件夹中。
  2. 修改代码中数据路径。
  3. 运行代码,训练和评估模型。

代码改进:

  • 图像特征提取: 可以使用更高级的图像特征提取方法,例如卷积神经网络,来获得更丰富的特征信息。
  • 边连接关系: 可以根据实际情况使用不同的边连接关系,例如基于节点距离或相似度的连接关系。
  • 模型架构: 可以尝试使用更复杂的图神经网络模型,例如 GAT 或 GraphSAGE,以提高分类性能。

总结:

本项目使用图神经网络对42个时刻的图数据进行了分类,并取得了不错的效果。通过进一步改进代码,可以获得更优的分类性能。

基于图神经网络的42时刻图数据分类

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

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