基于PyTorch和PyTorch Geometric的GCN多标签分类模型构建

本文介绍如何使用PyTorch和PyTorch Geometric库构建一个包含CNN和GCN的多标签分类模型。该模型以多个时刻的图像数据作为输入,利用GCN网络学习节点之间的关系,并进行多标签分类。

数据集构成:

  • 特征文件为'C:/Users/jh/Desktop/data/input/images/i.png_j.png',包含42个时刻的图片数据,每个时刻有37张图片(即37个节点)。
  • 图片尺寸为40 x 40。
  • 节点之间的连接关系相同,存储在'C:/Users/jh/Desktop/data/input/edges_L.csv'文件中,没有表头,第一列为源节点,第二列为目标节点,共有61条无向边。
  • 每个节点有8个标签,储存在'C:/Users/jh/Desktop/data/input/labels/i_j.txt'文本文件中,8个标签用空格隔开。
  • 将每个图的前30个节点图片颜色特征加入训练掩码,后7个节点图片颜色特征加入验证掩码。

模型架构:

  • CNN模型: 用于对输入的图像进行特征提取和降维,将每个节点的图片信息转化为一个特征向量。
  • GCN模型: 利用图神经网络学习节点之间的关系,并根据节点的特征和连接关系进行分类。

训练过程:

  • 模型以一个时刻的37个节点的数据作为输入。
  • CNN模型对输入数据进行降维,输出一个大小为(N, D)的二维张量,其中N表示节点数等于37,D表示每个节点的特征维度。
  • 将降维后的特征输入GCN模型,进行多标签分类。
  • 使用BCEWithLogitsLoss作为损失函数,Adam作为优化器,进行模型训练和验证。

代码示例:

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch_geometric.data import Data, DataLoader
from torch_geometric.nn import GCNConv
from torchvision.transforms import ToTensor
from PIL import Image
import pandas as pd
import numpy as np

# 设置文件路径
image_path = 'C:/Users/jh/Desktop/data/input/images/i.png_j.png'
label_path = 'C:/Users/jh/Desktop/data/input/labels/i_j.txt'
edge_path = 'C:/Users/jh/Desktop/data/input/edges_L.csv'

# 设置图像尺寸和节点数
image_size = (40, 40)
num_nodes = 37

# 图像预处理函数
transform = ToTensor()

# 定义图像数据集类
class ImageDataset(torch.utils.data.Dataset):
    def __init__(self, image_path, label_path, image_size):
        self.image_path = image_path
        self.label_path = label_path
        self.image_size = image_size
        self.data = self.load_data()

    def __getitem__(self, index):
        image = Image.open(self.image_path.format(index + 1))
        image = image.resize(self.image_size)
        image = transform(image)
        
        with open(self.label_path.format(index + 1), 'r') as f:
            labels = [float(label) for label in f.read().split()]
        labels = torch.tensor(labels)
        
        return image, labels

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

    def load_data(self):
        data = []
        for i in range(42):
            image = Image.open(self.image_path.format(i + 1))
            image = image.resize(self.image_size)
            image = transform(image)
            
            with open(self.label_path.format(i + 1), 'r') as f:
                labels = [float(label) for label in f.read().split()]
            labels = torch.tensor(labels)
            
            data.append((image, labels))
        
        return data

# 加载边的连接关系
edge_data = pd.read_csv(edge_path, header=None)
edge_index = torch.tensor(edge_data.values, dtype=torch.long).t().contiguous()

# 创建图数据
data_list = []
for i in range(42):
    x = torch.zeros(num_nodes, image_size[0] * image_size[1])
    y = torch.zeros(num_nodes, 8)
    
    for j in range(num_nodes):
        image, labels = ImageDataset(image_path, label_path, image_size).__getitem__(i * num_nodes + j)
        x[j] = image.view(-1)
        y[j] = labels
    
    data_list.append(Data(x=x, edge_index=edge_index, y=y))

# 创建数据加载器
train_loader = DataLoader(data_list[:30], batch_size=1, shuffle=True)
val_loader = DataLoader(data_list[30:], batch_size=1, shuffle=False)

# 定义CNN模型
class CNN(nn.Module):
    def __init__(self, image_size, num_nodes):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.fc = nn.Linear(64 * (image_size[0]//4) * (image_size[1]//4), num_nodes * 8)
        
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# 定义GCN模型
class GCN(nn.Module):
    def __init__(self, num_nodes, num_features, num_classes):
        super(GCN, self).__init__()
        self.conv1 = GCNConv(num_features, 16)
        self.conv2 = GCNConv(16, num_classes)
        
    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)
        return x

# 初始化CNN和GCN模型
cnn_model = CNN(image_size, num_nodes)
gcn_model = GCN(num_nodes, image_size[0] * image_size[1], 8)

# 定义损失函数和优化器
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(list(cnn_model.parameters()) + list(gcn_model.parameters()), lr=0.01)

# 训练函数
def train(cnn_model, gcn_model, train_loader, optimizer):
    cnn_model.train()
    gcn_model.train()
    
    for data in train_loader:
        x, y = data.x, data.y
        optimizer.zero_grad()
        
        x = cnn_model(x.unsqueeze(1))
        x = x.view(-1, num_nodes, 8)
        x = x.transpose(0, 1)
        x = x.reshape(num_nodes, -1)
        
        output = gcn_model(x, edge_index)
        loss = criterion(output, y)
        loss.backward()
        optimizer.step()

# 验证函数
def validate(cnn_model, gcn_model, val_loader):
    cnn_model.eval()
    gcn_model.eval()
    
    with torch.no_grad():
        total_loss = 0
        for data in val_loader:
            x, y = data.x, data.y
            
            x = cnn_model(x.unsqueeze(1))
            x = x.view(-1, num_nodes, 8)
            x = x.transpose(0, 1)
            x = x.reshape(num_nodes, -1)
            
            output = gcn_model(x, edge_index)
            loss = criterion(output, y)
            total_loss += loss.item()
        
        return total_loss / len(val_loader)

# 训练和验证模型
for epoch in range(10):
    train(cnn_model, gcn_model, train_loader, optimizer)
    val_loss = validate(cnn_model, gcn_model, val_loader)
    print('Epoch: {:02d}, Validation Loss: {:.4f}'.format(epoch+1, val_loss))

注意:

  • 请根据您的实际路径和需求进行适当的修改。
  • 此代码仅供参考,您可能需要根据您的具体任务和数据进行调整和改进。
  • 本代码中使用PyTorch和PyTorch Geometric库,您需要确保已安装这些库。
基于PyTorch和PyTorch Geometric的GCN多标签分类模型构建

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

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