定义无向加权图,并使用Pytorch_geometric实现图卷积

首先定义无向边并定义边的权重

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

a = torch.LongTensor([0, 0, 1, 1, 2, 2, 3, 4])
b= torch.LongTensor([0, 1, 2, 3, 1, 5, 1, 4])


num_A = 5
# 让b重新编号
b = b+num_A

# [源节点,目标节点]
first_c = torch.cat([a, b], dim=-1)
# [目标节点,源节点]
second_c = torch.cat([b, a], dim=-1)
# 拼接变为双向边
edge_index = torch.stack([first_c, second_c], dim=0)
# 因为双向边,把权重的维度要和边的个数匹配
rat = [0.5, 0.8, 1.0, 0.9, 0.7, 0.6,0.2,0.4]
ratings = torch.tensor(rat+rat, dtype=torch.float)
# 定义图
# edge_weight是权重特征,每条边有一个值,即[1,3]
# 如果想要为每条边定义多个特征,例如[[1,2],[2,3]]可以使用edge_attr
graph_data = Data(x=None, edge_index=edge_index,edge_weight=ratings)

print(graph_data.is_undirected())

 最后使用图卷积

class GraphConvNet(nn.Module):
    def __init__(self, graph_data):
        super(GraphConvNet, self).__init__()
 
        self.A_embeddings = nn.Embedding(5, 20)
        self.B_embeddings = nn.Embedding(6, 20)


        # 定义图卷积层
        self.conv1 = GCNConv(20, 20 // 2)
        self.conv2 = GCNConv(20 // 2, 20)
        self.norm = torch.nn.BatchNorm1d(20 // 2)
        self.data = graph_data
        self.data.x = (torch.cat([self.A_embeddings.weight, self.B_embeddings.weight], dim=0))

    def forward(self):
        x, edge_index,edge_weight = self.data.x, self.data.edge_index,self.data.edge_weight
   
        x = self.conv1(x, edge_index,edge_weight.view(-1))
        x = self.norm(x)
        x = torch.relu(x)
        x = F.dropout(x)
        x = self.conv2(x, edge_index,edge_weight)
        A_embedded = x[:5]
        B_embedded = x[5:]

        return A_embedded, B_embedded

gcnmodel = GraphConvNet(graph_data)
A_emb,B_emb = gcnmodel.forward()

你可能感兴趣的:(小记,pytorch,人工智能,python)