如何在LibTorch中释放GPU张量内存并保留CPU数据？

你遇到的核心问题是：torch::Tensor::to()方法会创建一个新的张量实例，而原来的GPU张量仍然被const auto& floatInput和const auto& target持有引用，所以GPU内存不会被立刻回收。

解决步骤 #

• 调用to(torch::kCPU)时，将结果保存到新的CPU张量变量中，确保后续能使用这些CPU数据
• 手动解除对GPU张量的引用，让PyTorch的CUDA缓存分配器可以回收对应GPU内存
• 调用emptyCache()的前提是已经没有张量引用目标内存，否则无法生效

修改后的代码示例 #

for(auto &batch: *dataloader){
    // 创建GPU张量，避免用const引用延长其生命周期
    auto floatInputGPU = batch.data.to(torch::kCUDA);
    auto targetGPU = batch.target.to(torch::kCUDA);
    
    size_t freeByte, totalByte;
    cudaMemGetInfo(&freeByte, &totalByte);
    cout << freeByte << " / " << totalByte <<  endl;

    // 保存CPU张量，保留所需数据
    auto floatInputCPU = floatInputGPU.to(torch::kCPU);
    auto targetCPU = targetGPU.to(torch::kCPU);
    
    // 显式重置GPU张量，释放内存引用
    floatInputGPU.reset();
    targetGPU.reset();
    
    // 清空缓存，此时GPU内存会被释放
    c10::cuda::CUDACachingAllocator::emptyCache();
    cudaMemGetInfo(&freeByte, &totalByte);
    cout << freeByte << " / " << totalByte << endl;
    
    // 后续可使用floatInputCPU和targetCPU处理数据
}

关键说明 #

• 原代码中const auto&的绑定会延长GPU张量的生命周期，即便生成了CPU张量，原GPU张量依然存在，内存无法释放
• reset()方法会销毁张量内部的存储，让CUDA分配器可以回收对应的内存块
• 若需要在循环外保留CPU数据，可将floatInputCPU和targetCPU存入外部容器（比如std::vector）

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