python通过导入自定义dll实现显存释放

博主研究了很多python的代码，发现都是无法释放显存的，或者说是释放显存不彻底。为此实现了自定义dll库，由python调用实现了释放显存。支持pytorch、tensorflow、onnxruntime等cuda运行环境，其运行效果与ai框架无关，由cuda C++ API决定。其中，自定义dll库导出给python使用可以参考vs2019导出动态链接库（dll）给其他vs项目及python代码使用_万里鹏程转瞬至的博客-CSDN博客通过vs可以导出动态链接库（dll文件）给其他c++项目、c#项目、python项目使用。本案例实现将vs项目导出为动态链接库，给c++项目与python项目使用。涉及全局变量、函数、自定义类的导出。项目创建完成后会得到以下结构， 可以将核心代码写在dllmain.cpp里面（原先的内容可以不用管），头文件信息写入在pch.h里面以下内容可以全部拷贝到pch.h中（博主的代码涉及到了cuda，所以需要配置以下cuda，cuda的配置可以参考libtorch显存管理示例_万里鹏程转瞬至的博客-CSDN博客，各https://hpg123.blog.csdn.net/article/details/125396626

下面代码中所用到的dll文件可以在以下链接下载，也可以参考上述链接自行实现。

python释放cuda缓存库-深度学习文档类资源-CSDN下载博主自行实现的动态链接库，通过python导入后可以实现释放显存，与ai框架无关。支持pytorch更多下载资源、学习资料请访问CSDN下载频道.https://download.csdn.net/download/a486259/85725926

 1、导入dll库

这些要调用cuda，因此电脑需配置好cuda环境。同时，注意按照自己的实际情况修改dll文件的路径。此外，代码用到了pynvml进行显存查询，如果电脑没有安装请用 pip install pynvml  进行安装。

import ctypes
import os
import pynvml
 
pynvml.nvmlInit()
def get_cuda_msg(tag=""):
    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
    meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
    print(tag, ", used:",meminfo.used / 1024**2,"Mib","free:",meminfo.free / 1024**2,"Mib")
    
#os.environ['Path']+=r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.1\bin'
#python3.7版本以上使用下列代码添加依赖项dll的路径
os.add_dll_directory(r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.1\bin')
lib = ctypes.cdll.LoadLibrary(os.getcwd()+ "/dll_export.dll")
#win32api.FreeLibrary(libc._handle)   #发现程序运行结束时无法正常退出dll，需要显式释放dll
#lib.reset_cuda()

2、搭建pytorch模型

这里所搭建的模型只是一个demo，没有任何的实际意义，各位可以使用自己的模型和框架。

import torch
import torch.nn as nn
import torch.nn.functional as F
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(6, 256, kernel_size=3, stride=1, padding=0)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=0)
        self.conv3 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=0)
        self.conv4 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=0)
        self.conv5 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=0)
        self.conv61 = nn.Conv2d(512, 3, kernel_size=3, stride=1, padding=0)
        self.conv62 = nn.Conv2d(512, 3, kernel_size=3, stride=1, padding=0)
        self.global_pooling = nn.AdaptiveAvgPool2d(output_size=(1, 1))
    #假设x1与x2的通道数均为3
    def forward(self, x1, x2):
        x = torch.cat([x1, x2], dim=1)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        x1= self.conv61(x)
        x2= self.conv62(x)
        x1 = self.global_pooling(x1).view(-1)
        x2 = self.global_pooling(x2).view(-1)
        x1 = F.softmax(x1,dim=0)
        x2 = F.softmax(x2,dim=0)
        return x1,x2

3、显存释放管理

代码及运行结果如下所示，包含模型初始化、变量初始化、forword后、del变化后、显存释放后（reset_cuda）5个状态前后的显存变化。可以看到调用lib.reset_cuda后显存基本上已经恢复。

get_cuda_msg("Context start")
model=MyModel().cuda()
get_cuda_msg("\n init model")
x1=torch.ones((10,3,256,256)).cuda()
x2=torch.ones((10,3,256,256)).cuda()
get_cuda_msg("\n init input")
model.forward(x1,x2)
get_cuda_msg("\n finish forward")
del x1,x2,model
get_cuda_msg("\n delete variable and model")
lib.reset_cuda()
get_cuda_msg("\n lib.reset_cuda()")