LibTorch 中的固定内存

Question

我可能遗漏了一些非常基本的东西，但我在文档或网上找不到任何解释

我正在尝试将 GPU at::Tensor 复制到 CPU 上的固定张量，但是一旦我复制它，CPU 张量就不再固定。我假设它只是创建了一个新的 GPU 张量副本并分配它，但如果是这样的话，你如何复制到一个预先分配的固定内存？

我的测试代码：

    at::Tensor gpu = at::randn({1025,1025}, device(at::kCUDA));
    at::Tensor pinned = at::empty(gpu.sizes(), device(at::kCPU).pinned_memory(true));
    std::cout << "Is Pinned: " << std::boolalpha << pinned.is_pinned() << std::endl;
    pinned = gpu.to(at::kCPU);
    std::cout << "Is Pinned: " << std::boolalpha << pinned.is_pinned() << std::endl;

输出为

Is Pinned: true
Is Pinned: false

torch:: 而不是 at::

也会发生这种情况

使用从源代码编译的 LibTorch 1.5.0 在 Ubuntu 16.04 上测试

Answer 1

我找到了一个方法，那就是使用 copy_ 函数

...
//pinned = gpu.to(torch::kCPU, true);
gpu.copy_(pinned);
std::cout << "Is Pinned: " << std::boolalpha << pinned.is_pinned() << std::endl;

这输出

Is Pinned: true
Is Pinned: true

我想这是有道理的，因为 to 函数 returns 是张量而不是操纵。虽然我希望 to 的某些变体允许它。

哦，好吧，它是这样工作的。

LibTorch 中的固定内存

Pinned memory in LibTorch

pytorch

tensor

libtorch