基于CNN-GCN的图像特征提取和图神经网络模型训练
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 = []
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('RGB').resize((224, 224), resample=Image.BILINEAR)
img_tensor = transforms.ToTensor()(img)
self.features.append(img_tensor)
# 读取标签
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)
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
network_id = idx // 37 # Calculate the network index
node_id = idx % 37 # Calculate the node index within the network
# Determine the edges for the current network
network_edges = self.edges[
(self.edges[0].isin(range(network_id * 37, (network_id + 1) * 37)))
& (self.edges[1].isin(range(network_id * 37, (network_id + 1) * 37)))
]
# Construct the edge_index tensor for undirected edges
source_nodes = network_edges[0].values
target_nodes = network_edges[1].values
edge_index = torch.tensor([source_nodes, target_nodes], dtype=torch.long)
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
if node_id < 30:
train_mask[node_id + network_id * 37] = 1
else:
val_mask[node_id + network_id * 37] = 1
data = Data(x=x, edge_index=edge_index, 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__()
self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
self.fc = nn.Linear(32 * 56 * 56, 40 * 40 * 3)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 32 * 56 * 56)
x = self.fc(x)
return x
# 创建GCN模型
class GCN(torch.nn.Module):
def __init__(self, num_node_features, num_classes):
super(GCN, self).__init__()
self.cnn = CNN() # 添加CNN模型
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)
x = F.softmax(x, dim=1)
x = x.view(-1, 37, -1) # 调整输出维度
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()
features = model.cnn(data.x) # 提取图像特征
data.x = features.view(features.size(0), -1)
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)
features = model.cnn(data.x) # 提取图像特征
data.x = features.view(features.size(0), -1)
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')
cnn_model = CNN().to(device)
model = GCN(num_node_features=40 * 40 * 3, 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=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}')
注意事项：
- 在导入库时，每个导入语句应位于单独的行。
- 在使用
pd.read_csv加载csv文件时，应指定列名，以便更好地理解数据。 - 在筛选边的连接关系时，使用逻辑运算符
&替代and。 - 在训练模型时，将图像特征提取放在GCN模型之前，以便在GCN模型中使用。
- 在验证模型时，使用
model.eval()指定模型为评估模式。 - 在计算准确率时，将
data.val_mask中的val_mask替换为train_mask。
边的连接关系如下： 0 1 0 4 0 8 1 2 1 5 2 3 2 6 3 7 3 14 4 5 4 9 5 6 5 10 6 7 6 11 7 13 8 9 8 15 9 10 9 17 10 11 10 18 11 12 12 13 12 19 13 14 13 20 14 21 15 16 16 17 16 22 17 18 17 23 18 24 19 20 19 28 20 21 20 29 21 30 22 23 22 25 23 24 23 26 24 27 25 26 25 34 26 27 26 31 27 28 27 31 28 32 29 30 29 33 30 36 31 32 31 34 32 33 32 35 33 36 34 35 35 36，储存在“C:\Users\jh\Desktop\data\input\edges_L.csv”中
原文地址: https://www.cveoy.top/t/topic/pdxV 著作权归作者所有。请勿转载和采集!