YOLOv7模型执行后GPU内存占用未下降的问题求助

核心问题分析 #

你每次调用mapImage函数都会重新加载一次YOLOv7模型，导致GPU上存在两个独立的模型实例，这是内存翻倍的主要原因。此外，未彻底清理推理产生的GPU张量、未正确使用梯度上下文管理也会加剧内存占用。

可行解决方案 #

1. 只加载一次模型（关键）

将模型加载逻辑移到mapImage函数外部，全局仅执行一次，避免重复创建模型实例占用GPU内存：

import matplotlib.pyplot as plt
import torch
import cv2
from torchvision import transforms
import numpy as np
from utils.datasets import letterbox
from utils.general import non_max_suppression_kpt
from utils.plots import output_to_keypoint
import uuid

# --------------------------
# 全局加载模型，仅执行一次
# --------------------------
c_or_g_pu = 'cuda:0' if torch.cuda.is_available() else 'cpu'
device = torch.device(c_or_g_pu)
weigths = torch.load(f'mods/yolov7-w6-pose.pt', map_location=device)
model = weigths['model']
_ = model.float().eval()
model.to(device)

def mapImage(image_path: str, id: str = None):
    if id is None: 
        id = str(uuid.uuid4())

    image = cv2.imread(image_path)
    image_path = f"temp/{id}.png"
    cv2.imwrite(image_path, image)

    output_path = f'temp/{id}_result.png'

    # --------------------------
    # 移除原函数内的模型加载代码
    # --------------------------

    image = cv2.imread(image_path)
    image = letterbox(image, 960, stride=64, auto=True)[0]
    image_ = image.copy()
    image = transforms.ToTensor()(image)
    image = torch.tensor(np.array([image.numpy()]))

    image = image.to(device)

    # --------------------------
    # 用torch.no_grad()包裹整个推理流程
    # --------------------------
    with torch.no_grad():
        output, _ = model(image)
        output = non_max_suppression_kpt(output, 0.25, 0.65, nc=model.yaml['nc'], nkpt=model.yaml['nkpt'], kpt_label=True)
        output = output_to_keypoint(output)

    markers = ImageMarkers(output)

    nimg = image[0].permute(1, 2, 0) * 255
    nimg = nimg.cpu().numpy().astype(np.uint8)
    nimg = cv2.cvtColor(nimg, cv2.COLOR_RGB2BGR)

    plt.imsave(output_path, nimg)
    mskShape = nimg.shape

    # --------------------------
    # 清理GPU张量并释放缓存
    # --------------------------
    del image, output
    torch.cuda.empty_cache()

    ...

2. 彻底清理GPU内存

• 使用del删除不再需要的GPU张量（如image、output），确保没有引用残留
• 调用torch.cuda.empty_cache()强制释放PyTorch持有的未使用GPU内存缓存
• 注意：torch.cuda.reset_peak_memory_stats()仅用于重置内存统计数据，不能释放实际内存，所以之前的尝试无效

3. 额外优化建议

• 避免重复读写图片：原代码中读取图片后又保存再读取，可直接使用初始加载的image变量，减少IO操作和内存临时占用
• 确保所有中间张量处理后及时移回CPU或删除，比如nimg已经转成CPU numpy数组，后续无需保留GPU上的image张量

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