Inference consumes all RAM but GPU RAM is not used

[valentinitnelav]

This is what I have tried:

import os
import torch
import yolov5

# Get the path to the project folder
prj_path = os.getcwd()
prj_path = prj_path.split('sequence_predictions')[0]
prj_path

# Load images
dir_img = "~/img_roi_crop" # custom path to directory with images on which I want to run inference
img_names = [name for name in os.listdir(dir_img) if name.endswith('.jpg')]
img_paths = [os.path.join(dir_img, name) for name in img_names]

# Load custom weights
weights_path = os.path.join(prj_path, "weights/paper_01/yolov5_s_best.pt") # some custom weights
model = yolov5.load(weights_path)

torch.cuda.is_available() 
# True

next(model.parameters()).device 
# device(type='cuda', index=0)

next(model.parameters()).is_cuda
# True

# set model parameters
model.conf = 0.1  # NMS confidence threshold
model.iou = 0.5  # NMS IoU threshold
model.max_det = 300  # maximum number of detections per image

# This runs out of the 64 Gb of RAM pretty fast. Why?
results = model(img_paths, size=1280, augment=True)
# I also tried , but I also run out of RAM
# results = model(img_paths, size=640, augment=False)

# Cannot execute code, session has been disposed. Please try restarting the Kernel.
# The Kernel crashed while executing code in the the current cell or a previous cell. 
# Please review the code in the cell(s) to identify a possible cause of the failure. 

I used watch -d -n 0.5 nvidia-smi to monitor the GPU usage and it looks like it's RAM is not being used actually. The GPU Memory Usage stays constantly at 1490MiB while the general 64 OS RAM is consumed in few minutes (as shown by htop command):

+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 530.30.02              Driver Version: 530.30.02    CUDA Version: 12.1     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                  Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf            Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA RTX A6000                On | 00000000:81:00.0 Off |                  Off |
| 30%   33C    P8               19W / 300W|   1493MiB / 49140MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A    687151      C   /bin/python3                               1490MiB |
+---------------------------------------------------------------------------------------+

Output of htop:

Could you help me understand where the bug is?