如何获取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页面中