将10张图片进行节点分类每张图片包含20个节点从节点0到节点19每张图66条边储存在CUsersjhDesktopdatarawedgescsv中每行为一条边第1列为源节点第2列为目标节点;1个节点包含2个特征所有图片节点的第一个特征储存在CUsersjhDesktopdatarawfeatures1csv中第二个特征储存在CUsersjhDesktopdatarawfeatures2csv中其中
import os
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch_geometric.data import Data
from torch_geometric.nn import GCNConv
from torch_geometric.data import DataLoader
class My_dataset(torch.utils.data.Dataset):
def __init__(self, root_dir):
self.root_dir = root_dir
self.edges_path = os.path.join(root_dir, 'edges.csv')
self.features1_path = os.path.join(root_dir, 'features1.csv')
self.features2_path = os.path.join(root_dir, 'features2.csv')
self.label_path = os.path.join(root_dir, 'label.csv')
self.edges = pd.read_csv(self.edges_path, header=None).values
self.features1 = pd.read_csv(self.features1_path, header=None).values
self.features2 = pd.read_csv(self.features2_path, header=None).values
self.labels = pd.read_csv(self.label_path, header=None).values
def __len__(self):
return len(self.features1)
def __getitem__(self, idx):
edge_index = torch.tensor(self.edges, dtype=torch.long).t().contiguous()
x = torch.cat((torch.tensor(self.features1, dtype=torch.float),
torch.tensor(self.features2, dtype=torch.float)), dim=1)
y = torch.tensor(self.labels, dtype=torch.float).squeeze()
data = Data(x=x, edge_index=edge_index, y=y)
return data
dataset = My_dataset("C:/Users/jh/Desktop/data/raw")
class GCN(nn.Module):
def __init__(self):
super(GCN, self).__init__()
self.conv1 = GCNConv(2, 16)
self.conv2 = GCNConv(16, 2)
def forward(self, data):
x, edge_index = data.x, data.edge_index
x = self.conv1(x, edge_index)
x = torch.relu(x)
x = self.conv2(x, edge_index)
return x
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GCN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)
train_loader = DataLoader(dataset[:8], batch_size=1, shuffle=True)
val_loader = DataLoader(dataset[8:], batch_size=1, shuffle=False)
def train(model, loader):
model.train()
for data in loader:
data = data.to(device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, data.y)
loss.backward()
optimizer.step()
def evaluate(model, loader):
model.eval()
correct = 0
for data in loader:
data = data.to(device)
output = model(data)
pred = output.max(dim=1)[1]
correct += int((pred == data.y).sum())
acc = correct / len(loader.dataset)
return acc
for epoch in range(10):
train(model, train_loader)
train_acc = evaluate(model, train_loader)
val_acc = evaluate(model, val_loader)
print(f'Epoch: {epoch+1}, Train Acc: {train_acc:.4f}, Val Acc: {val_acc:.4f}')
在上述代码中,首先定义了自定义数据集My_dataset,读取并存储了边、特征和标签的数据。然后定义了GCN模型类GCN,包含两个GCNConv层,并实现了前向传播函数。接着使用训练集和验证集的数据进行模型的训练和验证,使用Adam优化器和交叉熵损失函数进行模型训练。最后,使用循环迭代10次,输出每次迭代的训练集准确率和验证集准确率。
原文地址: https://www.cveoy.top/t/topic/i6gh 著作权归作者所有。请勿转载和采集!