add_self_loops
时间: 2024-05-12 16:12:24 浏览: 266
`add_self_loops`是用于在PyTorch几何库中的图形上添加自环的函数。自环是指从节点到它自身的边。在图神经网络中,为节点添加自环是一个常用的操作,可以增加节点的自我影响。下面是一个使用示例:
```python
import torch
from torch_geometric.utils import add_self_loops
from torch_sparse import SparseTensor
# 创建一个简单的图
edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype=torch.long)
edge_weight = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float)
adj = SparseTensor(row=edge_index, col=edge_index, value=edge_weight)
# 添加自环
adj = adj.to_symmetric()
adj, _ = add_self_loops(adj)
# 显示添加自环后生成的图
print(adj.coalesce())
```
在这个例子中,我们首先创建了一个具有四个节点和四条边的图。然后我们通过把图转化为对称矩阵,并使用`add_self_loops`函数来添加自环。最后我们输出了添加自环后生成的图。
相关问题
def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False, add_self_loops=True, flow="source_to_target", dtype=None): fill_value = 2. if improved else 1. if isinstance(edge_index, SparseTensor): assert flow in ["source_to_target"] adj_t = edge_index if not adj_t.has_value(): adj_t = adj_t.fill_value(1., dtype=dtype) if add_self_loops: adj_t = fill_diag(adj_t, fill_value) deg = sparsesum(adj_t, dim=1) deg_inv_sqrt = deg.pow_(-0.5) deg_inv_sqrt.masked_fill_(deg_inv_sqrt == float('inf'), 0.) adj_t = mul(adj_t, deg_inv_sqrt.view(-1, 1)) adj_t = mul(adj_t, deg_inv_sqrt.view(1, -1)) return adj_t else: assert flow in ["source_to_target", "target_to_source"] num_nodes = maybe_num_nodes(edge_index, num_nodes) if edge_weight is None: edge_weight = torch.ones((edge_index.size(1), ), dtype=dtype, device=edge_index.device) if add_self_loops: edge_index, tmp_edge_weight = add_remaining_self_loops( edge_index, edge_weight, fill_value, num_nodes) assert tmp_edge_weight is not None edge_weight = tmp_edge_weight row, col = edge_index[0], edge_index[1] idx = col if flow == "source_to_target" else row deg = scatter_add(edge_weight, idx, dim=0, dim_size=num_nodes) deg_inv_sqrt = deg.pow_(-0.5) deg_inv_sqrt.masked_fill_(deg_inv_sqrt == float('inf'), 0) return edge_index, deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col]
这段代码是一个用于计算GCN中归一化邻接矩阵的函数。下面是对代码的解读:
输入参数:
- edge_index:图的边索引,可以是一个包含两行的长为2的Tensor,表示边的起点和终点的索引。也可以是一个SparseTensor对象,表示稀疏的边索引。
- edge_weight:边的权重,可选参数,默认为None。如果不提供权重,则默认为全1.
- num_nodes:节点的数量,可选参数,默认为None。如果不提供数量,则通过edge_index推断得到。
- improved:布尔值,表示是否采用改进的归一化方式,默认为False。
- add_self_loops:布尔值,表示是否添加自环,默认为True。
- flow:字符串,表示信息传播的方向,默认为"source_to_target",即从源节点到目标节点。
函数内部逻辑:
1. 根据是否是SparseTensor对象,判断是稀疏还是稠密的边索引。
2. 对于稀疏边索引,首先判断信息传播方向是否为"source_to_target",然后进行按行归一化的处理。具体操作包括:
- 填充值的选择根据是否采用改进的方式而定。
- 如果边索引没有值,则填充为1。
- 如果需要添加自环,则对稀疏边索引进行填充对角线操作。
- 计算每个节点的度矩阵。
- 计算度矩阵的逆平方根。
- 对邻接矩阵进行按行归一化。
3. 对于稠密边索引,首先根据信息传播方向进行处理。具体操作包括:
- 如果需要添加自环,则对边索引和权重进行添加自环操作。
- 根据信息传播方向,选择相应的索引进行节点度计算。
- 计算度矩阵的逆平方根。
- 对边索引进行按元素归一化。
最终返回归一化后的邻接矩阵或者边索引。
def forward(self, x: Tensor, edge_index: Adj, edge_attr: OptTensor = None) -> Tensor: """""" if isinstance(edge_index, SparseTensor): edge_attr = edge_index.storage.value() if edge_attr is not None: edge_attr = self.mlp(edge_attr).squeeze(-1) if isinstance(edge_index, SparseTensor): edge_index = edge_index.set_value(edge_attr, layout='coo') if self.normalize: if isinstance(edge_index, Tensor): edge_index, edge_attr = gcn_norm(edge_index, edge_attr, x.size(self.node_dim), False, self.add_self_loops) elif isinstance(edge_index, SparseTensor): edge_index = gcn_norm(edge_index, None, x.size(self.node_dim), False, self.add_self_loops) x = self.lin(x) # propagate_type: (x: Tensor, edge_weight: OptTensor) out = self.propagate(edge_index, x=x, edge_weight=edge_attr, size=None) if self.bias is not None: out += self.bias return out
这是一个神经网络模型的前向传播函数。它接受输入张量 x 和边的索引 edge_index,以及可选的边属性 edge_attr。函数首先检查 edge_index 是否为稀疏张量类型,如果是,则将 edge_attr 设置为 edge_index 的值。然后,如果 edge_attr 不为空,则通过多层感知机(mlp)对其进行处理,并将维度压缩为一维。接下来,如果 edge_index 是稀疏张量类型,则使用 gcn_norm 函数对 edge_index 和 edge_attr 进行归一化处理。归一化过程中会使用 x 的维度信息和是否添加自环的标志位。然后,通过一个线性层对输入 x 进行变换。最后,调用 propagate 函数进行信息传递,并将结果加上偏置项(如果存在)。最终返回输出结果 out。
阅读全文