评价本页

张量索引 API#

在 PyTorch C++ API 中对张量进行索引与 Python API 非常相似。所有索引类型,如 None / ... / integer / boolean / slice / tensor,在 C++ API 中都可用,这使得将 Python 索引代码转换为 C++ 变得非常简单。主要区别在于,C++ API 没有使用类似于 Python API 语法的 [] 运算符,而是使用以下索引方法:

  • torch::Tensor::index (链接)

  • torch::Tensor::index_put_ (链接)

同样重要的是要注意,None / Ellipsis / Slice 等索引类型位于 torch::indexing 命名空间中,建议在使用这些索引类型之前加上 using namespace torch::indexing,以便于方便使用。

以下是一些将 Python 索引代码转换为 C++ 的示例:

获取器#

Python

C++ (假设使用了 using namespace torch::indexing)

tensor[None]

tensor.index({None})

tensor[Ellipsis, ...]

tensor.index({Ellipsis, "..."})

tensor[1, 2]

tensor.index({1, 2})

tensor[True, False]

tensor.index({true, false})

tensor[1::2]

tensor.index({Slice(1, None, 2)})

tensor[torch.tensor([1, 2])]

tensor.index({torch::tensor({1, 2})})

tensor[..., 0, True, 1::2, torch.tensor([1, 2])]

tensor.index({"...", 0, true, Slice(1, None, 2), torch::tensor({1, 2})})

设置器#

Python

C++ (假设使用了 using namespace torch::indexing)

tensor[None] = 1

tensor.index_put_({None}, 1)

tensor[Ellipsis, ...] = 1

tensor.index_put_({Ellipsis, "..."}, 1)

tensor[1, 2] = 1

tensor.index_put_({1, 2}, 1)

tensor[True, False] = 1

tensor.index_put_({true, false}, 1)

tensor[1::2] = 1

tensor.index_put_({Slice(1, None, 2)}, 1)

tensor[torch.tensor([1, 2])] = 1

tensor.index_put_({torch::tensor({1, 2})}, 1)

tensor[..., 0, True, 1::2, torch.tensor([1, 2])] = 1

tensor.index_put_({"...", 0, true, Slice(1, None, 2), torch::tensor({1, 2})}, 1)

Python/C++ 索引类型之间的转换#

Python 和 C++ 索引类型之间的一一对应转换如下:

Python

C++ (假设使用了 using namespace torch::indexing)

None

None

Ellipsis

Ellipsis

...

"..."

123

123

True

true

False

false

: or ::

Slice() or Slice(None, None) or Slice(None, None, None)

1: or 1::

Slice(1, None) or Slice(1, None, None)

:3 or :3:

Slice(None, 3) or Slice(None, 3, None)

::2

Slice(None, None, 2)

1:3

Slice(1, 3)

1::2

Slice(1, None, 2)

:3:2

Slice(None, 3, 2)

1:3:2

Slice(1, 3, 2)

torch.tensor([1, 2])

torch::tensor({1, 2})