如何获得 torch::Tensor 形状

Question

如果我们<<一个torch::Tensor

#include <torch/script.h>
int main()
{
    torch::Tensor input_torch = torch::zeros({2, 3, 4});
    std::cout << input_torch << std::endl;
    return 0;
}

我们看到了

(1,.,.) = 
  0  0  0  0
  0  0  0  0
  0  0  0  0

(2,.,.) = 
  0  0  0  0
  0  0  0  0
  0  0  0  0
[ CPUFloatType{2,3,4} ]

如何得到张量形状（即2,3,4）？我在 https://pytorch.org/cppdocs/api/classat_1_1_tensor.html?highlight=tensor 中搜索了一个 API 电话，但没有找到。而且我搜索了operator<<重载代码，也没找到。

Answer 1

你可以使用torch::sizes()方法

IntArrayRef sizes()

相当于python中的形状。此外，您可以通过调用 torch::size(dim) 访问给定 ax（维度）的特定大小。这两个功能都在您链接的 API 页面中

Answer 2

嗯，我一直在使用 torch::_shape_as_tensor(tensor)，它给了你另一个张量对象

Answer 3

对我有用的是：

#include <torch/script.h>
int main()
{
    torch::Tensor input_torch = torch::zeros({2, 3, 4});

    std::cout << "dim 0: " << input_torch.sizes()[0] << std::endl;
    std::cout << "dim 1: " << input_torch.sizes()[1] << std::endl;
    std::cout << "dim 2: " << input_torch.sizes()[2] << std::endl;

    assert(input_torch.sizes()[0]==2);
    assert(input_torch.sizes()[1]==3);
    assert(input_torch.sizes()[2]==4);
    
    return 0;
}

平台：

libtorch 1.11.0
CUDA 11.3

如何获得 torch::Tensor 形状

How to get torch::Tensor shape

c++

pytorch

libtorch