search

facebookresearch / segment-anything

The repository provides code for running inference with the SegmentAnything Model (SAM), links for downloading the trained model checkpoints, and example notebooks that show how to use the model.
Apache License 2.0
47.02k stars 5.57k forks source link

Segment anything doesn't use CUDA/GPU #557

Open LucasColas opened 1 year ago

LucasColas commented 1 year ago

Hello,

It seenms segment anything doesn't use GPU. Whenever I run a boiler code like this one : 

from segment_anything import SamPredictor, sam_model_registry
sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth")
predictor = SamPredictor(sam)
predictor.set_image(image)
masks, scores, logits= predictor.predict(point_coords=input_point,
    point_labels=input_label,
    multimask_output=True,
    )

I open my task manager. CPU is used highly more than GPU. GPU is barely used (1%). However for a few seconds my gpu use peaks at 100% (but the code doesn't take only a few seconds to be executed).

It’s replying true for torch.cuda.is_available(), but overall training speed and task manager’s graph seems torch can’t utilize GPU well.

heyoeyo commented 1 year ago

Most likely you just need to move the model over to the GPU in order to get it to work properly. Something like:

sam.to(device = "cuda")

Which you'd want to add just after creating the sam variable, before creating the predictor. Normally you also need to do this for the input data (i.e. the image in this case), but the set_image(...) function should handle that for you.

theFilipko commented 1 year ago

It works! :)