How to release memory space which occupied by segment-anything？

[forestbat]

I run this code in jupyterlab：

import cv2
import matplotlib.pyplot as plt
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
import numpy as np
import gc

def show_anns(anns):
    if len(anns) == 0:
        return
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()
    ax.set_autoscale_on(False)
    polygons = []
    color = []
    for ann in sorted_anns:
        m = ann['segmentation']
        img = np.ones((m.shape[0], m.shape[1], 3))
        color_mask = np.random.random((1, 3)).tolist()[0]
        for i in range(3):
            img[:,:,i] = color_mask[i]
        ax.imshow(np.dstack((img, m*0.35)))

sam = sam_model_registry["default"](checkpoint="VIT_H SAM Model/sam_vit_h_4b8939.pth")
mask_generator = SamAutomaticMaskGenerator(sam)
image = cv2.imread('Untitled Folder/292282 sample.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
sam.to(device='cuda')
masks = mask_generator.generate(image)
print(len(masks))
print(masks[0].keys())
plt.figure(figsize=(20,20))
plt.imshow(image)
show_anns(masks)
plt.axis('off')
plt.show() 
del(masks)
gc.collect()

Before I run it, memory cost is about 200MB, and when it completes, memory cost is about 3.4GB, even if I close the notebook or re-run this program, those memory won't be released. So how to solve the problem?

[VladKatsman]

try it out: torch.cuda.empty_cache()

[forestbat]

try it out: torch.cuda.empty_cache()

It has no effect.

[popay97]

try halving the input image resolution:

# Resize the image to a smaller resolution
scale_factor = 0.5
new_width = int(img.shape[1] * scale_factor)
new_height = int(img.shape[0] * scale_factor)
image = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_AREA)

# Convert the image from BGR to RGB color space
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

and/or running on CPU if you have more than 12gb+ ram, when I tested it rn on the CPU to see the time it would take it consumed roughly 12.5Gb on average.

[forestbat]

try halving the input image resolution:

# Resize the image to a smaller resolution
scale_factor = 0.5
new_width = int(img.shape[1] * scale_factor)
new_height = int(img.shape[0] * scale_factor)
image = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_AREA)

# Convert the image from BGR to RGB color space
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

and/or running on CPU if you have more than 12gb+ ram, when I tested it rn on the CPU to see the time it would take it consumed roughly 12.5Gb on average.

What I want to talk is not "the program costs too much memory", is "the program can't release memory completely", but your advice is still useful.

[YHSI5358]

I'm currently encountering this problem as well. Do you have any good solutions?

[Leotiv-Vibs]

Yeah, same problem. Please help to solve this issue!

[AlinGrigoras]

The solution is to use a multiprocessing approach, so you should move the image processing task to a separate process. By doing so, you isolate the memory usage of the task from the main process. This ensures that any memory consumed during processing is released back to the system when the process terminates, effectively preventing the memory leak