Memory leaks 循环中的特征提取似乎会导致pytorch中的内存泄漏

我花了相当多的时间尝试调试一些pytorch代码，我创建了一个最小的示例，目的是帮助更好地理解可能存在的问题

我已经删除了代码中与问题无关的所有必要部分，因此剩余的代码从功能角度看没有多大意义，但它仍然显示了我面临的错误

我正在处理的整个任务是在一个循环中，循环的每一步都是计算图像的嵌入并将其添加到存储它的变量中。它有效地聚合了它（而不是连接，因此大小保持不变）。我不期望迭代次数会迫使数据类型溢出，我在这里和代码中都没有看到这种情况

• 我添加了多个度量来评估我正在使用的张量的大小，以确保它们不会在内存占用中增长
• 我正在检查总体GPU内存使用情况，以验证导致最终运行时错误的问题：CUDA内存不足。

我的环境如下：

 - python 3.6.2
 - Pytorch 1.4.0
 - Cudatoolkit 10.0
 - Driver version 410.78
 - GPU: Nvidia GeForce GT 1030  (2GB VRAM)   
(though I've replicated this experiment with the same result on a Titan RTX with 24GB,
same pytorch version and cuda toolkit and driver, it only goes out of memory further in the loop).

完成下面的代码。我已经将2行标记为罪魁祸首，因为删除它们会消除问题，不过显然我需要找到一种方法来执行它们而不会出现内存问题。任何帮助都将不胜感激！您可以尝试使用任何名为“source_image.bmp”的图像来复制该问题

import torch
from PIL import Image
import torchvision
from torchvision import transforms
from pynvml import nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo, nvmlInit
import sys
import os
os.environ["CUDA_VISIBLE_DEVICES"]='0'      # this is necessary on my system to allow the environment to recognize my nvidia GPU for some reason
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'    # to debug by having all CUDA functions executed in place
torch.set_default_tensor_type('torch.cuda.FloatTensor')

# Preprocess image
tfms = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224), 
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),])
img = tfms(Image.open('source_image.bmp')).unsqueeze(0).cuda()

model = torchvision.models.resnet50(pretrained=True).cuda()
model.eval()    # we put the model in evaluation mode, to prevent storage of gradient which might accumulate

nvmlInit()
h = nvmlDeviceGetHandleByIndex(0)
info = nvmlDeviceGetMemoryInfo(h)
print(f'Total available memory   : {info.total / 1000000000}')

feature_extractor = torch.nn.Sequential(*list(model.children())[:-1])
orig_embedding = feature_extractor(img)

embedding_depth = 2048

mem0 = 0

embedding = torch.zeros(2048, img.shape[2], img.shape[3]) #, dtype=torch.float)

patch_size=[4,4]
patch_stride=[2,2]
patch_value=0.0

# Here, we iterate over the patch placement, defined at the top left location
for row in range(img.shape[2]-1):
    for col in range(img.shape[3]-1):
        print("######################################################")        
                
        ######################################################
        # Isolated line, culprit 1 of the GPU memory leak
        ######################################################
        patched_embedding = feature_extractor(img)
        
        delta_embedding = (patched_embedding - orig_embedding).view(-1, 1, 1)
        
        ######################################################
        # Isolated line, culprit 2 of the GPU memory leak
        ######################################################
        embedding[:,row:row+1,col:col+1] = torch.add(embedding[:,row:row+1,col:col+1], delta_embedding)

        print("img size:\t\t", img.element_size() * img.nelement())
        print("patched_embedding size:\t", patched_embedding.element_size() * patched_embedding.nelement())
        print("delta_embedding size:\t", delta_embedding.element_size() * delta_embedding.nelement())
        print("Embedding size:\t\t", embedding.element_size() * embedding.nelement())

        del patched_embedding, delta_embedding
        torch.cuda.empty_cache()
        
        info = nvmlDeviceGetMemoryInfo(h)
        print("\nMem usage increase:\t", info.used / 1000000000 - mem0)
        mem0 = info.used / 1000000000
        print(f'Free:\t\t\t {(info.total - info.used) / 1000000000}')

print("Done.")

---

加载模型后立即将其添加到代码中

模型参数（）中参数的

：
param.requires_grad=False

---

从

加载模型后立即将其添加到代码中

模型参数（）中参数的

：
param.requires_grad=False

从