【图神经网络之神器】torch_geometric

GCN/GraphSAGE/GAT代码

  • 导包

    import torch
    import torch.nn.functional as F
    from torch_geometric.nn import GCNConv, SAGEConv, GATConv
    from torch_geometric.datasets import Planetoid
    
  • 导入数据集

    dataset = Planetoid(root='./tmp/Cora', name='Cora')
    print(dataset.num_node_features)	# 节点的特征数 1433
    print(dataset.num_classes)			# 节点的类别数 7
    data = dataset[0]
    print(data.y)	# 节点对应的类别 
    print(data.x)	# 节点特征矩阵  [2708, 1433]
    print(data.edge_index)  # 图的边关系	[2, 10556]
    print(data.train_mask)  # 为true的位置,代表是训练集 140
    print(data.val_mask)	# 为true的位置,代表是验证集 500
    print(data.test_mask)	# 为true的位置,代表是测试集 1000
    

1. GCN模型

  • 构建

    class GCN_Net(torch.nn.Module):
        def __init__(self, features, hidden, classes):
            super(GCN_Net, self).__init__()
            self.conv1 = GCNConv(features, hidden)
            self.conv2 = GCNConv(hidden, classes)
    
        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_Net(dataset.num_node_features, 16, dataset.num_classes).to(device)
    data = dataset[0].to(device)	
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    
    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.test_mask].eq(data.y[data.test_mask]).sum()
    acc = int(correct) / int(data.test_mask.sum())
    print('GCN:', acc)
    

2. GraphSAGE模型

  • 构建

    class GraphSAGE_Net(torch.nn.Module):
        def __init__(self, features, hidden, classes):
            super(GraphSAGE_Net, self).__init__()
            self.sage1 = SAGEConv(features, hidden)
            self.sage2 = SAGEConv(hidden, classes)
    
        def forward(self, data):
            x, edge_index = data.x, data.edge_index
            x = self.sage1(x, edge_index)
            x = F.relu(x)
            x = F.dropout(x, training=self.training)
            x = self.sage2(x, edge_index)
            return F.log_softmax(x, dim=1)
    
  • 训练和测试

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = GraphSAGE_Ne.to(device)t(dataset.num_node_features, 16, dataset.num_classes).to(device)
    data = dataset[0].to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    
    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.test_mask].eq(data.y[data.test_mask]).sum()
    acc = int(correct) / int(data.test_mask.sum())
    print('GraphSAGE', acc)
    

3. GAT模型

  • 构建

    class GAT_Net(torch.nn.Module):
        def __init__(self, features, hidden, classes, heads=1):
            super(GAT_Net, self).__init__()
            self.gat1 = GATConv(features, hidden, heads=heads)
            self.gat2 = GATConv(hidden * heads, classes)
    
        def forward(self, data):
            x, edge_index = data.x, data.edge_index
            x = self.gat1(x, edge_index)
            x = F.relu(x)
            x = F.dropout(x, training=self.training)
            x = self.gat2(x, edge_index)
            return F.log_softmax(x, dim=1)
    
  • 训练和测试

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = GAT_Net(dataset.num_node_features, 16, dataset.num_classes, heads=4).to(device)
    data = dataset[0].to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    
    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.test_mask].eq(data.y[data.test_mask]).sum()
    acc = int(correct) / int(data.test_mask.sum())
    print('GAT', acc)
    

你可能感兴趣的:(知识图谱,神经网络,深度学习,机器学习,pytorch,图计算)