在使用libtorch進行部署時，會面臨顯存不夠用的情況。因此需要對顯存的利用進行管理，對此研究libtorch的api，嘗試進行顯存管理。libtorch運行程序時，顯存占用可以分為3塊：模型參數占用顯存、輸入輸出tensor占用顯存、模型forword過程臨時變量占用顯存。

使用cudaFree(tensor.data_ptr())可以釋放掉tensor所占用的顯存，也可以使用該函數釋放掉模型參數所占用的顯存。使用CUDACachingAllocator::emptyCache函數可以釋放掉模型在forword過程中的一些顯存。

為了探究模型部署時各階段的顯存占用和管理，進行以下代碼測試。

1、環境配置

進行顯存管理時，需要配置cuda，cuda配置可以參考圖1。此外，還需要在鏈接器-》輸入-》附加依賴項中，配置cudnn.lib;cublas.lib;cudart.lib;。

圖1 CUDA配置

libtorch的配置可以參考pytorch 4 libtorch配置使用實錄（支持cuda的調用）_萬里鵬程轉瞬至的博客-CSDN博客

?2、庫導入和基本函數實現

下列代碼主要實現cuda和libtorch的導入，和cuda使用查詢的函數

#include <torch/script.h> #include <torch/torch.h> #include <c10/cuda/CUDAStream.h> #include <ATen/cuda/CUDAEvent.h>#include <iostream> #include <memory> #include <string>#include <cuda_runtime_api.h> using namespace std;static void print_cuda_use( ) {size_t free_byte;size_t total_byte;cudaError_t cuda_status = cudaMemGetInfo(&free_byte, &total_byte);if (cudaSuccess != cuda_status) {printf("Error: cudaMemGetInfo fails, %s \n", cudaGetErrorString(cuda_status));exit(1);}double free_db = (double)free_byte;double total_db = (double)total_byte;double used_db_1 = (total_db - free_db) / 1024.0 / 1024.0;std::cout << "Now used GPU memory " << used_db_1 << " MB\n"; }

3、libtorch過程顯存管理

代碼中的d_in_out.pt參考自以下鏈接，但略有不同（模型參數更多了，博主把模型的參數增加了100倍，其實就是把每一個kernel的filer_num增大了）pytorch 6 libtorch部署多輸入輸出模型（支持batch）_萬里鵬程轉瞬至的博客-CSDN博客_pytorch多輸入1、pytorch下構建多輸入輸出模型下面構建一個簡潔的多輸入輸出模型import torchimport torch.nn as nnimport torch.nn.functional as Fclass MyModel(nn.Module): def __init__(self): super(MyModel, self).__init__() self.conv1 = nn.Conv2d(6, 16, kernel_size=1, stridehttps://blog.csdn.net/a486259/article/details/121680988

?代碼中的注釋詳細說明了每一個操作是否有效和其影響?

int main() {string path = "d_in_out.pt";//設置不需要存儲梯度信息at::NoGradGuard nograd;int gpu_id = 0;//加載模型torch::jit::Module model = torch::jit::load(path);model.to(at::kCUDA);model.eval();//設置評價模式std::cout << "加載模型后的顯存占用\n";print_cuda_use();//構建雙輸入數據//對于單輸入模型只需要push_back一次at::Tensor x1_tensor = torch::ones({ 1,3,512,512 }).to(at::kCUDA);at::Tensor x2_tensor = torch::ones({ 1,3,512,512 }).to(at::kCUDA);at::Tensor result1,result2;std::cout << "\n初始化tensor的顯存占用\n";print_cuda_use();std::cout << "\n循環運行5次\n";for (int i = 0;i < 5;i++) {//result=model.forward({ x1_tensor,x2_tensor }).toTensor();//one outauto out = model.forward({ x1_tensor,x2_tensor });auto tpl = out.toTuple();//out.toTensorList();result1 = tpl->elements()[0].toTensor();result2 = tpl->elements()[1].toTensor();print_cuda_use();}std::cout << "\n釋放tensor所占用的顯存（幫助不大）\n";cudaFree(x1_tensor.data_ptr());cudaFree(x1_tensor.data_ptr());cudaFree(result1.data_ptr());cudaFree(result2.data_ptr());print_cuda_use();std::cout << "\nCUDACachingAllocator::emptyCache （有點效果）\n";c10::cuda::CUDACachingAllocator::emptyCache();print_cuda_use();std::cout << "\n釋放模型參數占用的顯存（無意義）\n";for (auto p : model.parameters())cudaFree(p.data_ptr());//接下來不能使用cuda，導致model.to(at::kCUDA);報錯print_cuda_use();//torch::jit::Module::Module::freeze(model);std::cout << "\n調用模型的析構函數（無效）\n";model.~Module();//對顯存變化沒有實際幫助print_cuda_use();std::cout << "\n重置cuda狀態（無效）\n";c10::cuda::CUDACachingAllocator::init(gpu_id);c10::cuda::CUDACachingAllocator::resetAccumulatedStats(gpu_id);//對顯存變化沒有實際幫助c10::cuda::CUDACachingAllocator::resetPeakStats(gpu_id);print_cuda_use();std::cout << "\ncudaDeviceReset（有效，會導致后續模型無法使用cuda）\n";//需要配置cudacudaDeviceReset();//完全釋放GPU資源 接下來不能使用cuda，導致model.to(at::kCUDA);報錯 除非能重新初始化libtorch的cuda環境print_cuda_use();torch::cuda::synchronize();model = torch::jit::load(path);model.to(at::kCUDA);print_cuda_use();return 0; }

上述代碼的執行結果如下圖所示。?

最終結果表明，libtorch是無法有效的釋放顯存。需要需要做到顯存的極致管理，還得使用onnxruntime或tensorrt進行部署。

onnxruntime的部署可以參考pytorch 17 onnx多輸入多輸出模型在python與C++下用OnnxRuntime部署_萬里鵬程轉瞬至的博客-CSDN博客_onnx多輸入pytorch模型在轉換成onnx模型后可以明顯加速，此外模型在進行openvino部署時也需要將pytorch模型轉換為onnx格式。為此，以多輸入多輸出模型為例，記錄一下模型轉換及python下onnxruntime調用過程。并實現C++下多輸入多輸出模型的Onnxruntime的調用。一 、python下模型轉onnx與測試1.1、構建pytorch多輸入多輸出模型import torchimport torch.nn as nnimport torch.nn.functional https://hpg123.blog.csdn.net/article/details/122473889tensorrt的部署可以參考pytorch 29 onnx多輸入多輸出模型（動態尺寸）轉TensorRT模型并在python與C++下用TensorRT進行部署_萬里鵬程轉瞬至的博客-CSDN博客_onnx模型轉tensorrt實現將多輸入多輸出的onnx模型轉TensorRT的格式，并用TensorRT的python api進行調用，最后實現C++的調用，python結果與C++結果一模一樣。在此過程中實現了，動態尺寸，也就是可以在運行是按照數據情況動態調整模型的輸入數據的格式。在這里以pytorch的多輸入多輸出模型為例，從模型構建到轉化為onnx動態size、tensorrt安裝、onnx轉tensorRT動態size、tensorRT模型pytthon調用、tensorRT C++項目配置，最終到tensorRT模型C++https://hpg123.blog.csdn.net/article/details/125191219

總結

以上是生活随笔為你收集整理的libtorch显存管理示例的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。