import torch from torch_geometric.data import Data, Dataset from torch_geometric.nn import GCNConv import torch.nn.functional as F

class My_dataset(Dataset): def init(self, root, transform=None, pre_transform=None): super(My_dataset, self).init(root, transform, pre_transform)

@property
def raw_file_names(self):
    return ['edges.txt', 'features1.txt', 'features2.txt', 'labels.txt']

@property
def processed_file_names(self):
    return ['data.pt']

def process(self):
    # 读取边数据
    with open(self.raw_paths[0], 'r') as f:
        edges = f.read().split()
        edges = [int(edge) for edge in edges]

    # 读取特征数据
    with open(self.raw_paths[1], 'r') as f:
        features1 = f.read().split()
        features1 = [int(feature) for feature in features1]

    with open(self.raw_paths[2], 'r') as f:
        features2 = f.read().split()
        features2 = [int(feature) for feature in features2]

    # 读取标签数据
    with open(self.raw_paths[3], 'r') as f:
        labels = f.read().split()
        labels = [int(label) for label in labels]

    # 创建图数据
    edge_index = torch.tensor(edges, dtype=torch.long).view(2, -1)
    features = torch.stack([torch.tensor(features1), torch.tensor(features2)], dim=1)
    labels = torch.tensor(labels, dtype=torch.long)
    data = Data(x=features, edge_index=edge_index, y=labels)

    # 保存图数据
    torch.save(data, self.processed_paths[0])

def len(self):
    return 1

def get(self, idx):
    return torch.load(self.processed_paths[0])

dataset = My_dataset(root=r'C:\Users\jh\Desktop\data')

划分训练集和验证集

train_mask = torch.zeros(10, dtype=torch.bool) train_mask[:6] = 1 val_mask = torch.zeros(10, dtype=torch.bool) val_mask[6:] = 1

data = dataset[0] data.train_mask = train_mask data.val_mask = val_mask

定义GCN模型

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 = F.dropout(x, training=self.training)
    x = self.conv2(x, edge_index)
    return F.log_softmax(x, dim=1)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = GCN().to(device) data = data.to(device)

训练模型

optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4) model.train() for epoch in range(200): optimizer.zero_grad() out = model(data) loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask]) loss.backward() optimizer.step()

验证模型

model.eval() _, pred = model(data).max(dim=1) correct = pred[data.val_mask].eq(data.y[data.val_mask]).sum().item() accuracy = correct / data.val_mask.sum().item()

print('Loss: {:.4f}, Accuracy: {:.4f}'.format(loss.item(), accuracy))