时间戳

首页 > 解决方案 > LibTorch 中的固定内存

pytorch - LibTorch 中的固定内存

问题描述

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

我正在尝试将 GPU 复制到at::TensorCPU 上的固定张量，但是一旦我复制它，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 上测试

标签： pytorchtensorlibtorch

解决方案

我找到了一种方法，那就是使用该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允许它。

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

推荐阅读