将10张图片进行节点分类每张图片包含20个节点节点从0到19每张图数据66条边储存在CUsersjhDesktopdatarawedgescsv中每行为一条边一共有66行第1列为源节点第2列为目标节点;1个节点包含2个特征所有图片节点的第一个特征储存在CUsersjhDesktopdatarawfeatures1csv中第二个特征储存在CUsersjhDesktopdatarawfeatures2
下面是使用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库,需要提前导入相应的模块。
原文地址: https://www.cveoy.top/t/topic/i6cJ 著作权归作者所有。请勿转载和采集!