首先定义无向边并定义边的权重
- 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()