如何获得 torch::张量形状

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)。这两个函数都位于您链接的 API 页面中

Answer 2

对我有用的是：

#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