基于图神经网络的图片多标签分类模型

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.edges = pd.read_csv(os.path.join(root, 'edges_L.csv'), header=None)
        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')
        label_path = os.path.join(self.root, 'labels', f'{idx+1}.txt')

        # 读取图片特征
        image = self.read_image_features(image_path)

        # 读取标签
        labels = self.read_labels(label_path)

        # 读取边关系
        edge_index = torch.tensor(self.edges.values, dtype=torch.long).t().contiguous()

        # 创建Data对象
        data = Data(x=image, edge_index=edge_index, y=labels)

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

        return data

    def read_image_features(self, image_path):
        # 读取图片特征，此处使用示例代码中的随机特征生成方法，请根据实际情况进行修改
        return torch.randn(37, 8)

    def read_labels(self, label_path):
        # 读取标签
        with open(label_path, 'r') as f:
            labels = [list(map(int, line.strip().split())) for line in f]
        return torch.tensor(labels, dtype=torch.long)

# 定义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}')

修改说明：

修改了MyDataset类的实现，添加了读取图片特征和标签的方法，并根据实际情况进行修改。
修改了GCN模型的输入特征维度，根据实际情况进行修改。
修改了train_model和validate_model函数中的output和data.y的使用，使其适应多标签分类任务。
修改了数据集路径的设置，将根目录设置为"data/input"，并将图片存放在"data/input/images"目录下，标签存放在"data/input/labels"目录下。边的关系存放在"data/input/edges_L.csv"文件中。请根据实际情况进行修改。
在训练过程中，打印每个epoch的训练损失和验证准确率。