首先定义无向边并定义边的权重
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()
更多【python-定义无向加权图,并使用Pytorch_geometric实现图卷积】相关视频教程:www.yxfzedu.com