首先定义无向边并定义边的权重
import torch
import torch.nn as nn
from torch_geometric.nn import GCNConv
import torch.nn.functional as F
from torch_geometric.data import Dataa = 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_dataself.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_weightx = 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_embeddedgcnmodel = GraphConvNet(graph_data)
A_emb,B_emb = gcnmodel.forward()
】)
)
)
