feat: Add parallel multi-GPU inference for VGGT inference in demo_colmap.py
#445
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
meta-cla
bot
added
the
CLA Signed
Adds optional multi-GPU mode that distributes VGGT inference across multiple GPUs using torch.multiprocessing for true parallel processing. Key changes: - Add --multi_gpu flag to enable parallel multi-GPU mode - Add --gpu_ids to specify which GPUs to use (default: all available) - Add _worker_process() for multiprocessing worker - Add run_VGGT_multi_gpu() for parallel inference orchestration - Use shared memory for efficient tensor transfer between processes - Add 'spawn' start method for CUDA compatibility in __main__ Multi-GPU benefits: - Memory: Each GPU loads only its shard (~55 frames instead of 221) - Speed: Near-linear speedup with parallel processing - Example: 221 images on 4 GPUs = ~55 frames/GPU, ~42GB peak each Usage: # All available GPUs python demo_colmap.py --scene_dir=/path/to/scene --multi_gpu # Specific GPUs python demo_colmap.py --scene_dir=/path/to/scene --multi_gpu --gpu_ids=0,1 Note: BA and tracking still run on single GPU after multi-GPU inference.
xiaoya27
force-pushed
the
feat/multi-gpu-inference
branch
from
dbf9299 to
d00d1d3
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
Adds optional parallel multi-GPU mode for VGGT inference in
demo_colmap.pyusingtorch.multiprocessing. This enables processing large image sequences that would otherwise exceed single GPU memory.Motivation
VGGT inference on 221 images requires ~77GB VRAM, which exceeds the capacity of many GPUs. This PR allows distributing frames across multiple GPUs to:
Changes
--multi_gpuflag to enable parallel multi-GPU mode--gpu_idsto specify which GPUs to use (default: all available)_worker_process()for multiprocessing workersrun_VGGT_multi_gpu()for parallel inference orchestrationUsage