本文旨在深入解析 NumPy 数组与 PyTorch 张量在索引操作上的差异，特别是在使用形状为 (1,) 的数组或张量进行索引时。通过对比 NumPy 和 PyTorch 的行为，揭示其底层机制，并提供清晰的示例和解释，帮助读者更好地理解和应用这两种常用的科学计算库。
NumPy 索引行为
NumPy 提供了强大的索引功能，允许使用整数、切片、布尔数组甚至其他 NumPy 数组来访问和修改数组元素。当使用 NumPy 数组作为索引时，NumPy 会将其解释为一组需要提取的元素的索引。
例如：
import numpy as np x = np.arange(10) y = x[np.array([1])] print(y) # 输出: [1]
在这个例子中，np.array([1]) 被解释为一个索引数组，NumPy 返回一个包含索引 1 处的元素的新数组。
PyTorch 张量索引行为
PyTorch 张量也支持索引操作，但其行为在某些情况下与 NumPy 数组有所不同。当使用 PyTorch 张量作为索引时，如果该张量是一个只包含单个元素的整数张量，PyTorch 会尝试将其转换为一个普通的 Python 整数索引。
例如：
import torch as th x = np.arange(10) z = x[th.tensor([1])] print(z) # 输出: 1
在这个例子中，th.tensor([1]) 首先被转换为整数 1，然后 x[1] 返回索引 1 处的元素。
差异的根源
NumPy 和 PyTorch 在处理张量索引时的差异源于它们对张量的不同解释。NumPy 倾向于将张量索引视为一组索引，而 PyTorch 尝试将单元素整数张量转换为标量索引。
为了更深入地理解这种差异，我们可以查看 NumPy 的源代码。当 NumPy 遇到一个非整数索引时，它会尝试调用该对象的 __index__ 方法。如果该方法存在且成功返回一个整数，NumPy 将使用该整数作为索引。
if (PyLong_CheckExact(obj) || !PyArray_Check(obj)) {
// it calls PyNumber_Index() internally
npy_intp ind = PyArray_PyIntAsIntp(obj);
if (error_converting(ind)) {
PyErr_Clear();
}
else {
index_type |= HAS_INTEGER;
indices[curr_idx].object = NULL;
indices[curr_idx].value = ind;
indices[curr_idx].type = HAS_INTEGER;
used_ndim += 1;
new_ndim += 0;
curr_idx += 1;
continue;
}
}PyTorch 的张量类实现了 __index__ 方法，该方法仅对包含单个元素的整数张量有效。
>>> import torch >>> torch.tensor([1]).__index__() 1 >>> torch.tensor([1, 2]).__index__() Traceback (most recent call last): File "<stdin>", line 1, in <module> TypeError: only integer tensors of a single element can be converted to an index
总结与注意事项
- 当使用 NumPy 数组作为索引时，NumPy 将其解释为一组索引。
- 当使用 PyTorch 张量作为索引时，如果该张量是包含单个元素的整数张量，PyTorch 会尝试将其转换为整数索引。
- 了解 NumPy 和 PyTorch 在索引行为上的差异对于避免意外错误至关重要。
- 在需要使用张量作为索引时，请务必明确所需的行为，并根据需要进行类型转换或调整索引方式。
通过理解这些差异，可以更有效地利用 NumPy 和 PyTorch 的索引功能，并编写更健壮和可预测的代码。
