LibTorch中TensorOptions的使用

$$LibTorch中TensorOptions的使用$$

TensorOptions是用来配置Tensor属性用的

auto options =  torch::TensorOptions()
    .dtype(torch::kFloat32)
    .layout(torch::kStrided)
    .device(torch::kCUDA, 1)
    .requires_grad(true);
1
2
3
4
5

torch::Tensor tensor = torch::full({3, 4}, /*value=*/123, options);

assert(tensor.dtype() == torch::kFloat32);
assert(tensor.layout() == torch::kStrided);
assert(tensor.device().type() == torch::kCUDA); // or device().is_cuda()
assert(tensor.device().index() == 1);
assert(tensor.requires_grad());
1
2
3
4
5
6
7

示例：

#include 
using namespace std;
int main()
{

    auto options = torch::TensorOptions()
        .dtype(torch::kFloat32)
        .requires_grad(true);

    torch::Tensor w = torch::randn({ 1 }, options);

    assert(w.dtype() == torch::kFloat32);
    assert(w.requires_grad());

    cout << (w.dtype() == torch::kFloat32) << "  " << w.requires_grad() << endl;

	system("pause");
	return 0;
}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

参考官方链接

https://pytorch.org/cppdocs/notes/tensor_creation.html