import os import pandas as pd import torch import torch.nn as nn from torch_geometric.data import Data from torch_geometric.data import DataLoader from torch_geometric.nn import GCNConv import torch.nn.functional as F from torchvision import transforms from PIL import Image from sklearn.model_selection import train_test_split

定义CNN网络

class CNN(nn.Module): def init(self, in_channels, out_channels): super(CNN, self).init() self.conv1 = nn.Conv2d(in_channels, 16, kernel_size=3, stride=1, padding=1) self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0) self.conv2 = nn.Conv2d(16, out_channels, kernel_size=3, stride=1, padding=1)

def forward(self, x):
    x = F.relu(self.conv1(x))
    x = self.pool(x)
    x = F.relu(self.conv2(x))
    x = self.pool(x)
    return x

定义GCN模型

class GCN(nn.Module): def init(self, in_channels, out_channels): super(GCN, self).init() self.conv1 = GCNConv(in_channels, 128) self.dropout = nn.Dropout(p=0.5) self.conv2 = GCNConv(128, out_channels)

def forward(self, data):
    x, edge_index = data.x, data.edge_index
    x = F.relu(self.conv1(x, edge_index))
    x = self.dropout(x)
    x = self.conv2(x, edge_index)
    return x

读取边的关系数据

edges = pd.read_csv('C:\Users\jh\Desktop\data\input\edges_L.csv', header=None) edges = edges.values

读取节点特征数据

features = [] labels = [] for i in range(1, 43): for j in range(37): image_path = f'C:\Users\jh\Desktop\data\input\images{i}.png_{j}.png' image = Image.open(image_path).convert('RGB') transform = transforms.Compose([transforms.Resize((40, 40)), transforms.ToTensor()]) image_tensor = transform(image) features.append(image_tensor)

    labels_path = f'C:\Users\jh\Desktop\data\input\labels\{i}_{j}.txt'
    with open(labels_path, 'r') as file:
        label = [int(l) for l in file.readline().strip().split()]
        labels.extend(label)

x = torch.stack(features) x = x.view(-1, 3, 40, 40)

划分训练集和验证集的掩码

mask_train = torch.zeros(42, 37, dtype=torch.bool) mask_val = torch.zeros(42, 37, dtype=torch.bool) for i in range(42): mask_train[i, :30] = 1 mask_val[i, 30:] = 1 mask_train = mask_train.view(-1) mask_val = mask_val.view(-1)

创建图结构

edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous() data_list = [] for i in range(42): data = Data(x=x, edge_index=edge_index) data.mask_train = mask_train[i*37:(i+1)37] data.mask_val = mask_val[i37:(i+1)37] data.y = torch.tensor(labels[i37:(i+1)*37]) data_list.append(data)

创建CNN模型实例，降维至8维

cnn_model = CNN(in_channels=3, out_channels=8)

使用CNN模型对节点特征进行降维

with torch.no_grad(): cnn_output = [] for i in range(42): x_i = x[i*37:(i+1)*37].unsqueeze(1) x_i = x_i.squeeze(1) output_i = cnn_model(x_i) output_i = output_i.view(output_i.size(0), -1) cnn_output.append(output_i) cnn_output = torch.cat(cnn_output, dim=0)

将降维后的特征加入data对象

for i in range(42): data_list[i].x = cnn_output[i*37:(i+1)*37]

创建GCN模型实例

gcn_model = GCN(in_channels=800, out_channels=8)

定义多标签分类的损失函数和优化器

criterion = nn.BCEWithLogitsLoss() optimizer = torch.optim.Adam(gcn_model.parameters(), lr=0.0008)

训练模型

num_epochs = 5000 for epoch in range(num_epochs): gcn_model.train() total_loss = 0 for data in data_list: optimizer.zero_grad() out = gcn_model(data) labels_batch = F.one_hot(data.y, num_classes=8).float() loss = criterion(out, labels_batch) loss.backward() optimizer.step() total_loss += loss.item()

avg_loss = total_loss / len(data_list)
print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {avg_loss:.6f}')

在验证集上评估模型

gcn_model.eval() with torch.no_grad(): total_correct = 0 total_samples = 0 for data in data_list: out = gcn_model(data) labels_batch = F.one_hot(data.y, num_classes=8).float() predicted = (out > 0).float() total_correct += (predicted == labels_batch).sum().item() total_samples += labels_batch.size(0) * labels_batch.size(1)

accuracy = total_correct / total_samples
print(f'Validation Accuracy: {accuracy:.4f}')