pytorch - 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 函数返回一个张量而不是进行操作。尽管我希望 to 的某些变体能够允许它。

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