基于颜色特征的图神经网络模型:使用42个时间点的37个节点图片数据进行训练
基于颜色特征的图神经网络模型:使用42个时间点的37个节点图片数据进行训练
本项目使用42个时间点的37个节点图片数据训练图神经网络模型。每个时间点包含37张图片,每张图片代表一个节点,并具有1600维的颜色特征。该模型采用GCN架构,通过学习节点之间的连接关系和颜色特征,对节点进行分类。
数据描述:
- 数据路径:'C:\Users\jh\Desktop\data\images_block'。
- 图片大小:40×40。
- 图片命名格式:'i.png_j.png',其中'i'表示时刻,从1到42,'j'表示节点序号,从0到36。
- 每个时刻包含37张图片,代表37个节点。
- 边的连接关系在所有时刻相同。
模型架构:
- GCN(图卷积网络)
- 输入特征:1600维颜色特征
- 输出类别:8
代码实现:
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_flatten', f'{i}.txt_{j}.txt')
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]
x = x.view(37, -1) # 调整特征数据维度
# 定义图数据的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}')
新特征输入:
新的特征输入应该是提取上述图片的所有颜色特征。在给定的代码中,特征数据已经被加载为torch.Tensor类型的对象,每个对象表示一张图片的颜色特征。因此,可以直接将这些特征数据作为GCN模型的输入。在GCN模型中,输入特征的维度应该等于num_node_features参数的值,即1600。所以,可以将特征数据的维度调整为(37, 1600),然后将其作为模型的输入。
具体而言,可以在MyDataset类中的__getitem__方法中添加以下代码来提取特征数据并将其调整为合适的维度:
features = self.features[idx]
features = features.view(37, -1)
然后,在GCN模型的forward方法中,将调整后的特征数据作为模型的输入:
x = self.conv1(features, edge_index)
这样,新的特征数据就会作为GCN模型的输入,用于训练和验证。
原文地址: https://www.cveoy.top/t/topic/pbnc 著作权归作者所有。请勿转载和采集!