LibTorch中GPU张量复制到预分配固定内存后属性丢失问题咨询

我完全懂你碰到的这个坑——你明明预分配了带固定内存属性的CPU张量，结果用to()赋值后，这个属性直接没了。核心原因其实很简单：gpu.to(at::kCPU)会生成一个全新的普通CPU张量，你把这个新张量赋值给pinned变量时，直接覆盖了原来指向固定内存张量的引用，所以后面查is_pinned()自然返回false。

要搞定这个问题，关键是直接把GPU张量的数据复制到预分配好的固定内存张量里，而不是创建新张量替换它。这里有两种靠谱的方法：

方法一：用copy_原地复制 #

copy_方法支持跨设备的原地数据拷贝，完全不会改变目标张量的内存属性。修改你的代码如下：

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;
// 直接把GPU数据复制到已有的固定内存张量中
pinned.copy_(gpu);
std::cout << "Is Pinned: " << std::boolalpha << pinned.is_pinned() << std::endl;

执行后两次输出都会是true，因为我们只是把数据拷贝到预分配的固定内存里，没有替换pinned指向的原始张量对象。

方法二：用to()的out参数指定目标张量 #

你也可以在调用to()时，通过out参数指定要写入的目标张量，这样就不会生成新的张量：

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;
// 通过out参数指定目标固定内存张量
gpu.to(at::kCPU, /*non_blocking*/ false, /*copy*/ true, pinned);
std::cout << "Is Pinned: " << std::boolalpha << pinned.is_pinned() << std::endl;

这个方法和copy_效果一致，都是原地修改预分配的固定内存张量，保留其内存属性。

额外提示 #

• 如果想提升性能用异步复制，可以把non_blocking参数设为true，但前提是你的GPU支持异步传输，且后续操作不依赖这次复制的结果。
• 不管用torch::还是at::命名空间，这两种方法都适用，因为它们底层调用的是同一套API。

内容的提问来源于stack exchange，提问作者La bla bla