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 from torchvision import transforms from PIL import Image

加载数据并创建PyG数据集类：

class MyDataset(torch.utils.data.Dataset): def init(self, root, transform=None, pre_transform=None): self.edges = pd.read_csv(os.path.join(root, 'input', 'edges_L.csv'), header=None) self.transform = transform self.pre_transform = pre_transform

    # 读取特征和标签数据
    self.features = []
    self.labels = []
    self.cnn = CNN()
    for i in range(1, 43):
        for j in range(37):
            # 读取特征
            img_name = os.path.join(root, 'input', 'images', '{}.png_{}.png'.format(i, j))
            img = Image.open(img_name).convert('RGBA').resize((40, 40), resample=Image.BILINEAR)
            img_tensor = transforms.ToTensor()(img)
            feature = self.cnn(img_tensor.unsqueeze(0))
            self.features.append(feature)
            
            # 读取标签
            label_name = os.path.join(root, 'input', 'labels', '{}_{}.txt'.format(i, j))
            with open(label_name, 'r') as f:
                labels = [int(x) for x in f.readline().strip().split()]
            self.labels.append(labels)
    
    # 将特征调整维度为[batch_size, num_node_features, width, height]
    self.features = torch.stack(self.features, dim=0)
    self.labels = torch.tensor(self.labels)
    
    # Calculate the total number of nodes
    self.num_nodes = len(self.labels)
    
    # 添加边的关系
    self.edges = torch.tensor(self.edges.values, dtype=torch.long).t().contiguous() - 1
    
    self.cnn = CNN()  # 添加CNN模型

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

def __getitem__(self, idx):
    x = self.features[idx]
    y = self.labels[idx]
    
    # Define graph-wide train_mask and val_mask
    train_mask = torch.zeros(self.num_nodes, dtype=torch.bool)
    val_mask = torch.zeros(self.num_nodes, dtype=torch.bool)
    
    # Set train_mask for the first 30 nodes in each network, and val_mask for the last 7 nodes
    node_id = idx % 37
    network_id = idx // 37
    if node_id < 30 and (node_id + network_id * 37) < self.num_nodes:
        train_mask[node_id + network_id * 37] = 1
    elif (node_id + network_id * 37) < self.num_nodes:
        val_mask[node_id + network_id * 37] = 1
    
    data = Data(x=x, edge_index=self.edges, y=y, train_mask=train_mask, val_mask=val_mask)
    
    if self.pre_transform is not None:
        data = self.pre_transform(data)
    
    if self.transform is not None:
        data.x = self.transform(data.x)
    
    return data

定义CNN模型

class CNN(torch.nn.Module): def init(self): super(CNN, self).init() # Modify the first convolution layer to reduce spatial dimensions self.conv1 = nn.Conv2d(4, 16, kernel_size=(3, 3), stride=2, padding=1) self.conv2 = nn.Conv2d(16, 4, kernel_size=3, stride=2, padding=1)

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

创建GCN模型

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) self.cnn = CNN()

def forward(self, data):
    x, edge_index = data.x, data.edge_index
    x = self.cnn(x)  # 使用CNN模型提取图像特征
    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(dataloader, model, optimizer, device): model.train() total_loss = 0.0

for data in dataloader:
    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(dataloader)

def validate_model(dataloader, model, device): model.eval() correct = 0 total = 0

for data in dataloader:
    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': dataset = MyDataset(root="C:\Users\jh\Desktop\data") device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = GCN(num_node_features=4, 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=True)
val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False)

epochs = 1

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