基于图神经网络的图像节点分类模型

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

class MyDataset(torch.utils.data.Dataset): def init(self, root, transform=None, pre_transform=None): self.transform = transform self.pre_transform = pre_transform self.data_list = []

    for i in range(1, 43):  # 处理42张图片，编号从1到42
        for j in range(37):  # 每张图片37个节点
            # 加载特征值数据
            features_file = os.path.join(root, 'input', 'images_flatten', f'{i}.txt_{j}.txt')
            features = pd.read_csv(features_file, header=None, sep=' ', encoding='utf-8')
            x = torch.tensor(features.values, dtype=torch.float)

            # 加载标签数据
            labels_file = os.path.join(root, 'input', 'labels', f'{i}.txt_{j}.txt')
            labels = pd.read_csv(labels_file, header=None, sep=' ', encoding='utf-8')
            y = torch.tensor(labels.values, dtype=torch.long).squeeze()

            # 构建图数据
            edge_index_file = os.path.join(root, 'input', 'edges_P.csv')
            edge_index = pd.read_csv(edge_index_file, header=None, sep=',', encoding='utf-8')
            edge_index = torch.tensor(edge_index.values, dtype=torch.long).t().contiguous()

            # 创建 train_mask 和 val_mask
            train_mask = torch.zeros(y.size(0), dtype=torch.bool)
            val_mask = torch.zeros(y.size(0), dtype=torch.bool)
            train_mask[:30] = 1  # 前30个节点作为训练集
            val_mask[30:] = 1  # 后7个节点作为验证集

            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)
            self.data_list.append(data)

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

def __getitem__(self, idx):
    return self.data_list[idx]

class GCN(torch.nn.Module): def init(self, num_node_features, num_classes): super(GCN, self).init() self.conv1 = GCNConv(num_node_features, 8) self.conv2 = GCNConv(8, 16) self.conv3 = GCNConv(16, 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.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': root = r'C:\Users\jh\Desktop\data' dataset = MyDataset(root=root) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = GCN(num_node_features=1600, num_classes=8).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

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

epochs = 10
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}')