YOLOv7 model GPU memory usage doesn't go down after execution

Question

I have a YOLOv7 model that I run with Pytorch - it executes twice in the script. I've earlier executed it in an environment that only had a CPU, but I just switched to a new environment with a GPU.

The problem is after it has been executed, the GPU memory usage doesn't go down. After the first execution at this line: output, _ = model(image) the GPU usage goes up to 2498 MiB and when the model executes the second time, the GPU usage goes up to 4572 MiB. I have another model that also needs to be executed, but because the GPU memory usage is so high, the script crashes.

What can I do to "reset" the GPU memory usage, after running the YOLOv7 model?

I've tried to del output after this line:

with torch.no_grad():
    output = output_to_keypoint(output)

and also tried to torch.cuda.reset_peak_memory_stats(device=None), but nothing seems to work and I don't really know what to try otherwise.

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

def mapImage(image_path: str, id: str = None):
    if type(id) == 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'

    # YOLOv7

    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)

    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)
    output, _ = model(image)

    output = non_max_suppression_kpt(output, 0.25, 0.65, nc=model.yaml['nc'], nkpt=model.yaml['nkpt'], kpt_label=True)

    with torch.no_grad():
        output = output_to_keypoint(output)

    markers = ImageMarkers(output)

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

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

    ...

I hope you have some suggestions!

YOLOv7 model GPU memory usage doesn't go down after execution

0 Answers0