下面是使用torch_geometric库自定义数据集My_dataset类实现的代码:

import torch
from torch_geometric.data import Data, Dataset

class My_dataset(Dataset):
    def __init__(self, root):
        self.root = root
        self.edges = torch.tensor([], dtype=torch.long)
        self.features1 = torch.tensor([], dtype=torch.float)
        self.features2 = torch.tensor([], dtype=torch.float)
        self.labels = torch.tensor([], dtype=torch.long)
        
        # Load data from files
        edges_file = root + "/edges.csv"
        features1_file = root + "/features1.csv"
        features2_file = root + "/features2.csv"
        labels_file = root + "/label.csv"

        self.edges = torch.from_numpy(np.genfromtxt(edges_file, delimiter=",", dtype=np.int32))
        self.features1 = torch.from_numpy(np.genfromtxt(features1_file, delimiter=",", dtype=np.float32))
        self.features2 = torch.from_numpy(np.genfromtxt(features2_file, delimiter=",", dtype=np.float32))
        self.labels = torch.from_numpy(np.genfromtxt(labels_file, delimiter=",", dtype=np.int32))
        
    def __len__(self):
        return 10
    
    def __getitem__(self, idx):
        edge_index = self.edges[idx].t().contiguous()
        features = torch.stack([self.features1[idx], self.features2[idx]], dim=-1)
        labels = self.labels[idx]
        
        train_mask = torch.zeros(features.size(0), dtype=torch.bool)
        train_mask[:16] = 1
        val_mask = torch.zeros(features.size(0), dtype=torch.bool)
        val_mask[16:] = 1
        
        data = Data(edge_index=edge_index, x=features, y=labels, train_mask=train_mask, val_mask=val_mask)
        return data

# Create dataset
dataset = My_dataset("C:/Users/jh/Desktop/data/raw")

# GCN model
class GCN(torch.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 = F.relu(x)
        x = self.conv2(x, edge_index)
        return x

model = GCN()

# Training and validation loop
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

for i in range(len(dataset)):
    data = dataset[i].to(device)
    optimizer.zero_grad()
    out = model(data)
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    with torch.no_grad():
        pred = out.argmax(dim=1)
        correct_train = int(pred[data.train_mask].eq(data.y[data.train_mask]).sum().item())
        train_acc = correct_train / int(data.train_mask.sum())
        correct_val = int(pred[data.val_mask].eq(data.y[data.val_mask]).sum().item())
        val_acc = correct_val / int(data.val_mask.sum())
    print('Epoch: {:03d}, Train Acc: {:.4f}, Val Acc: {:.4f}'.format(i, train_acc, val_acc))

以上代码实现了自定义数据集My_dataset类,将数据加载到torch_geometric.data.Data对象中,并创建了一个简单的GCN模型,利用训练集train_mask进行模型训练,并使用验证集val_mask进行模型验证。循环遍历10张图片的图数据,并将其输入模型进行训练和验证。请注意,代码中使用的GCNConv和F.relu来自于torch_geometric库,需要提前导入相应的模块。

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

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

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