基于图片颜色特征的图神经网络模型训练

'C:\Users\jh\Desktop\data\images_block'路径下的所有图片都是40×40大小,名字格式为i.png_j.png,其中i表示时刻，i从1到42，j表示节点序号，j从0到36，共有42个时刻的图数据,而且边的连接关系相同,每个时刻都有37张图片,每张图片代表1个节点,提取上述图片的颜色特征,作为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.edges = pd.read_csv(os.path.join(root, 'edges_L.csv')) self.transform = transform self.pre_transform = pre_transform self.num_classes = 8 # 修改成8类标签 self.features = [] self.labels = [] for i in range(1, 43): for j in range(37): feature_path = os.path.join(root, 'images_block', f'{i}.png_{j}.png') label_path = os.path.join(root, 'labels', f'{i}_{j}.txt') # 修改标签文件名 features = pd.read_csv(feature_path, header=None, sep=' ') labels = pd.read_csv(label_path, header=None, sep=' ', encoding='ansi') self.features.append(torch.tensor(features.values, dtype=torch.float)) self.labels.append(torch.tensor(labels.values.squeeze(), dtype=torch.long)) # 修改标签数据类型为long型 def len(self): return len(self.features) def getitem(self, idx): edge_index = torch.tensor(self.edges.values, dtype=torch.long).t().contiguous() x = self.features[idx] y = self.labels[idx] # 定义图数据的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[:30] = 1 val_mask[30:] = 1 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) return data # 定义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) 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': 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=1600, num_classes=8).to(device) # 修改num_classes为8 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 = 2 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}')