【PYG】切片语法[:,]

在 Python 中，使用切片语法 [:,] 是一种方便的方式来选择多维数组或张量的特定部分。具体来说，这种语法在处理 NumPy 数组或 PyTorch 张量时非常有用。让我们详细解释一下为什么可以写成 [:, :] 以及这种语法的含义。

切片语法解释

:：表示选择数组或张量的整个范围。在一维数组或张量中，: 表示从头到尾选择所有元素。
[:, :]：对于二维数组或张量，第一部分 : 表示选择所有行，第二部分 : 表示选择所有列。

这种语法的效果是选择二维数组或张量的所有元素。

补充一点个人想法，通过代码验证

array[:, :2] 是 array[:, 0:2] 的简略写法
取出前两行array[0:2, :]和array[:2, :]是一样的

print(array[0:2, :])
[[0 1 2]
 [3 4 5]]

组合使用 print(array[0:1, :2])，得到[[0 1]]

示例代码

以下是使用 NumPy 和 PyTorch 的一些示例，展示如何使用这种切片语法：

使用 NumPy

import numpy as np

# 创建一个二维 NumPy 数组
array = np.array([[0, 1, 2],
                  [3, 4, 5],
                  [6, 7, 8]])

# 打印整个数组
print("Entire array:")
print(array)

# 使用 [:, :] 选择整个数组
print("Using [:, :] to select the entire array:")
print(array[:, :])

# 使用 [:, :2] 选择所有行的前两列
print("Using [:, :2] to select all rows and first two columns:")
print(array[:, :2])

程序输出结果

Entire array:
[[0 1 2]
 [3 4 5]
 [6 7 8]]
Using [:, :] to select the entire array:
[[0 1 2]
 [3 4 5]
 [6 7 8]]
Using [:, :2] to select all rows and first two columns:
[[0 1]
 [3 4]
 [6 7]]

使用 PyTorch

import torch

# 创建一个二维 PyTorch 张量
tensor = torch.tensor([[0, 1, 2],
                       [3, 4, 5],
                       [6, 7, 8]])

# 打印整个张量
print("Entire tensor:")
print(tensor)

# 使用 [:, :] 选择整个张量
print("Using [:, :] to select the entire tensor:")
print(tensor[:, :])

# 使用 [:, :2] 选择所有行的前两列
print("Using [:, :2] to select all rows and first two columns:")
print(tensor[:, :2])

具体到 `edge_index` 示例

在 PyTorch Geometric 中，edge_index 是一个形状为 [2, num_edges] 的张量，其中 num_edges 是图中的边数量。使用 [:, :] 语法表示选择整个 edge_index 张量。具体到 edge_index[:, :10]，表示选择前 10 列（即前 10 条边）。

# 示例 PyTorch Geometric 数据
edge_index = torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
                           [1, 2, 0, 4, 5, 3, 7, 6, 5, 8]])

# 选择前 10 列
print("First 10 edges (edge_index[:, :10]):")
print(edge_index[:, :10])