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次,输出每次迭代的训练集准确率和验证集准确率。

将10张图片进行节点分类每张图片包含20个节点从节点0到节点19每张图66条边储存在CUsersjhDesktopdatarawedgescsv中每行为一条边第1列为源节点第2列为目标节点;1个节点包含2个特征所有图片节点的第一个特征储存在CUsersjhDesktopdatarawfeatures1csv中第二个特征储存在CUsersjhDesktopdatarawfeatures2csv中其中

原文地址: https://www.cveoy.top/t/topic/i6gh 著作权归作者所有。请勿转载和采集!

免费AI点我,无需注册和登录