torch_geometric.utils.remove_self_loops()
函数的作用是从图中删除自环。
以下是使用torch_geometric.utils.remove_self_loops()
的示例代码:
import torch
from torch_geometric.utils import remove_self_loops
from torch_geometric.data import Data
# 定义一个简单的图
edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 1, 1]], dtype=torch.long)
data = Data(edge_index=edge_index)
print(data.edge_index) # tensor([[0, 1, 1, 2],[1, 0, 1, 1]])
# 删除自环
data.edge_index, _ = remove_self_loops(data.edge_index)
print(data.edge_index) # tensor([[0, 1, 2],[1, 0, 1]])
在上面的代码中,我们首先定义了一个简单的图,然后使用torch_geometric.utils.remove_self_loops()
函数删除自环。函数返回的第一个元素是删除自环后的边索引,第二个元素是包含自环的索引。由于我们不需要自环,因此将第二个元素忽略了。