facebookresearch / segment-anything-2

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

Does vos_inference.py support multi-GPU running? #228

Open lzkzls opened 4 weeks ago

lzkzls commented 4 weeks ago

Thanks for the open source. I run vos_inference.py with my own video dataset on a machine with four v100 graphics cards (each with 32G memory). I detected that the program is only running on one card. Does vos_inference.py support multi-GPU running? How should I modify the code?

ronghanghu commented 4 weeks ago

Hi @lzkzls, currently vos_inference.py only runs on a single GPU. However, since it mostly just loop over the video list, you may split your video list into multiple chunks and use the --video_list_file parameter in https://github.com/facebookresearch/segment-anything-2/blob/7e1596c0b6462eb1d1ba7e1492430fed95023598/tools/vos_inference.py#L231 to perform inference on a separate chunk of the dataset in each GPU and save their prediction PNG files into the output directory.