该代码使用 PyTorch Geometric 库实现一个图神经网络 (GNN) 模型,用于对图像进行分类。模型使用 GCNConv 层,并使用自定义数据集进行训练和验证。

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='latin-1')  # 添加 encoding 参数
                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='latin-1')  # 添加 encoding 参数
                y = torch.tensor(labels.values, dtype=torch.long).squeeze()

                # 构建图数据
                edge_index = torch.zeros((x.size(0), 2), dtype=torch.long)
                edge_index[:, 1] = torch.arange(x.size(0))
                edge_index = edge_index.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[:16] = 1  # 前16个节点作为训练集
                val_mask[16:] = 1  # 后4个节点作为验证集

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

代码说明:

  • MyDataset 类读取特征值和标签数据,构建图数据,并划分训练集和验证集。
  • GCN 类定义一个三层 GCN 模型,使用 GCNConv 层进行卷积操作。
  • train_model 函数进行模型训练,计算训练损失。
  • validate_model 函数进行模型验证,计算验证集准确率。
  • 主函数加载数据,创建模型和优化器,进行训练和验证。

需要注意的是,要根据实际情况修改文件路径和参数。

代码改进:

  1. 添加了 encoding='latin-1' 参数到 pd.read_csv() 函数,解决文件解码问题。
  2. 对代码进行了注释,解释了代码的功能。
  3. 使用了 r'C:\Users\jh\Desktop\data' 形式表示文件路径,避免了转义字符问题。

更多信息:

  • PyTorch Geometric 库:https://pytorch-geometric.readthedocs.io/en/latest/
  • 图神经网络 (GNN) 的概念:https://en.wikipedia.org/wiki/Graph_neural_network

希望本代码能够帮助你实现 GNN 模型并应用于图像分类任务。

图神经网络 (GNN) 在图像分类中的应用:基于 PyTorch Geometric 的实现

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

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