MiDaS深度图计算出现内存泄漏，求原因及解决方法

我使用MiDaS的DPT_BEiT_L_512模型计算深度图，目前遇到的问题是：遍历不同图片进行推理时，RAM占用会持续上升；但重复使用同一张图片（即使添加随机噪声）或仅读取图片不进行推理时，内存占用保持稳定。

完整代码 #

#!venv/bin/python3

from pathlib import Path
import psutil

import numpy as np
import torch
import cv2


def make_model():
    model_type = "DPT_BEiT_L_512"       # MiDaS v3.1 - Large
    midas = torch.hub.load("intel-isl/MiDaS", model_type)

    device = torch.device("cuda")
    midas.to(device)
    midas.eval()

    midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")

    transform = midas_transforms.dpt_transform

    return {"transform": transform,
            "device": device,
            "midas": midas
            }


def inference(cv_image, model):
    """Make the inference."""
    transform = model['transform']
    device = model["device"]
    midas = model["midas"]

    input_batch = transform(cv_image).to(device)

    with torch.no_grad():
        prediction = midas(input_batch)

        prediction = torch.nn.functional.interpolate(
            prediction.unsqueeze(1),
            size=cv_image.shape[:2],
            mode="bilinear",
            align_corners=False,
        ).squeeze()

    output = prediction.cpu().numpy()
    formatted = (output * 255 / np.max(output)).astype('uint8')

    return formatted


# Create Midas "DPT_BEiT_L_512" -  MiDaS v3.1 - Large
model = make_model()

image_dir = Path('.') / "all_images"

for image_file in image_dir.iterdir():
    ram_usage = psutil.virtual_memory()[2]
    print("image", ram_usage)
    cv_image = cv2.imread(str(image_file))
    _ = inference(cv_image, model)

简要流程 #

• 创建DPT_BEiT_L_512模型实例
• 定义inference函数处理单张图片的深度图计算
• 遍历all_images目录下的所有图片
• 用cv2.imread读取每张图片
• 调用inference计算深度图（未保留计算结果）

测试变体 #

变体1：始终使用同一图片 #

替换循环代码后，内存占用稳定无泄漏：

cv_image = cv2.imread("image.jpg")
for i in range(1, 100):
    ram_usage = psutil.virtual_memory()[2]
    print(i, ram_usage)
    _ = inference(cv_image, model)

变体2：不计算深度图 #

仅读取图片不调用推理函数，内存无泄漏，排除cv2.imread的问题：

for image_file in image_dir.iterdir():
    ram_usage = psutil.virtual_memory()[2]
    print("image", ram_usage)
    cv_image = cv2.imread(str(image_file))
    # _ = inference(cv_image, model)

变体3：同一图片添加随机噪声 #

给同一张图片每次添加不同噪声后推理，内存占用保持稳定：

cv_image = cv2.imread("image.jpg")
for i in range(1, 100):
    ram_usage = psutil.virtual_memory()[2]
    print(i, ram_usage)

    noise = np.random.randn(
        cv_image.shape[0], cv_image.shape[1], cv_image.shape[2]) * 20
    noisy_img = cv_image + noise
    noisy_img = np.clip(noisy_img, 0, 255)

    _ = inference(noisy_img, model)

问题 #

1. 内存泄漏的来源是什么？
2. 两次推理之间是否需要重置某些内容？

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