diff --git a/README.md b/README.md index 98ba182c..07e36734 100644 --- a/README.md +++ b/README.md @@ -54,7 +54,7 @@ This makes visdet significantly easier to install and deploy compared to the ori Without access full training logs (loss plots etc.), it can be impossible to know if you have your own implementation wrong or not. Ideally, eventually we integrate the docs, and the experiment results into the same one living documentation. We run hyperparameter search, you get the new best hyperparameters. -Goals of the repo: +Goals of the repo: - Even more open than your typical open-source project, logs available, roadmap available. - Emphasis on all of DevEx, educating new users, production deployments and research.2 diff --git a/docs/user-guide/inference.md b/docs/user-guide/inference.md index 55cba251..ac927790 100644 --- a/docs/user-guide/inference.md +++ b/docs/user-guide/inference.md @@ -54,6 +54,84 @@ for image_file in image_files: # Process result... ``` +### Multi-GPU Inference + +To run inference on multiple GPUs in a single process, pass multiple CUDA devices to `init_detector` and then infer a batch (a list) of images: + +```python +from visdet.apis import inference_detector, init_detector + +model = init_detector(config_file, checkpoint_file, device="cuda:0,1") +results = inference_detector(model, image_files) # list[DetDataSample] +``` + +## FiftyOne (Voxel51) Dataset + +If you use [FiftyOne](https://voxel51.com/fiftyone/) for dataset management and visualization, you can run `visdet` inference over a `fiftyone.Dataset` and attach the predictions back onto each sample. + +```python +import fiftyone as fo + +from visdet.apis import inference_detector, init_detector +from visdet.utils import detections_to_fiftyone + +# 1) Load your model +config_file = "configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py" +checkpoint_file = "checkpoints/faster_rcnn_r50_fpn_1x_coco.pth" +model = init_detector(config_file, checkpoint_file, device="cuda:0") + +# 2) Create (or load) a FiftyOne dataset +# You can also do: dataset = fo.load_dataset("my-dataset") +dataset = fo.Dataset.from_images_dir( + "path/to/images", + name="my-images", + overwrite=True, +) +dataset.compute_metadata() # populates sample.metadata.width/height + +classes = model.dataset_meta.get("classes", []) +score_thr = 0.3 + +# 3) Run inference and attach detections +for sample in dataset.iter_samples(progress=True): + data_sample = inference_detector(model, sample.filepath) + + pred = data_sample.pred_instances + pred = pred[pred.scores > score_thr] + + width = sample.metadata.width + height = sample.metadata.height + + dets = [] + for bbox, label_id, score in zip( + pred.bboxes.cpu().numpy(), + pred.labels.cpu().numpy(), + pred.scores.cpu().numpy(), + strict=False, + ): + x1, y1, x2, y2 = bbox.tolist() + + dets.append( + { + "label": classes[int(label_id)] if classes else str(int(label_id)), + "bounding_box": [ + x1 / width, + y1 / height, + (x2 - x1) / width, + (y2 - y1) / height, + ], + "confidence": float(score), + } + ) + + sample["predictions"] = detections_to_fiftyone(dets) + sample.save() + +# 4) Visualize in the FiftyOne App +session = fo.launch_app(dataset) +session.wait() +``` + ## Test Dataset Evaluate model on a test dataset: diff --git a/tests/test_runtime/test_inference_device_parsing.py b/tests/test_runtime/test_inference_device_parsing.py new file mode 100644 index 00000000..e87ff1c3 --- /dev/null +++ b/tests/test_runtime/test_inference_device_parsing.py @@ -0,0 +1,19 @@ +import pytest + +from visdet.apis import inference as inference_api + + +def test_parse_inference_devices_single_device(): + assert inference_api._parse_inference_devices("cpu") == ("cpu", None) + assert inference_api._parse_inference_devices("cuda:0") == ("cuda:0", None) + + +def test_parse_inference_devices_multi_device(): + assert inference_api._parse_inference_devices("cuda:0,1") == ("cuda:0", [0, 1]) + assert inference_api._parse_inference_devices(["cuda:0", "cuda:1"]) == ("cuda:0", [0, 1]) + assert inference_api._parse_inference_devices([0, 2, 3]) == ("cuda:0", [0, 2, 3]) + + +def test_parse_inference_devices_rejects_non_cuda_sequences(): + with pytest.raises(ValueError): + inference_api._parse_inference_devices(["cpu", "cuda:0"]) # type: ignore[arg-type] diff --git a/visdet/apis/inference.py b/visdet/apis/inference.py index 5e8e5f06..12f582af 100644 --- a/visdet/apis/inference.py +++ b/visdet/apis/inference.py @@ -16,7 +16,7 @@ # from visdet.cv.ops import RoIPool # Removed - eliminating C++ ops from visdet.cv.transforms import Compose from visdet.engine.config import Config -from visdet.engine.dataset import default_collate +from visdet.engine.dataset import default_collate, pseudo_collate from visdet.engine.model.utils import revert_sync_batchnorm from visdet.engine.registry import init_default_scope from visdet.engine.runner import load_checkpoint @@ -27,11 +27,42 @@ from visdet.utils import ConfigType, get_test_pipeline_cfg +def _parse_inference_devices(device: str | Sequence[str] | Sequence[int]) -> tuple[str, list[int] | None]: + if isinstance(device, str): + if device.startswith("cuda") and "," in device: + parts = [p.strip() for p in device.split(",") if p.strip()] + device_ids = [_device_to_id(p) for p in parts] + return f"cuda:{device_ids[0]}", device_ids + return device, None + + device_ids = [_device_to_id(d) for d in device] + return f"cuda:{device_ids[0]}", device_ids + + +def _device_to_id(device: str | int) -> int: + if isinstance(device, int): + return device + + if device.isdigit(): + return int(device) + + if device == "cuda": + return 0 + + if device.startswith("cuda:"): + return int(device.split(":", 1)[1]) + + raise ValueError( + "Multi-device inference only supports CUDA devices. " + f"Got device={device!r}; expected e.g. 'cuda:0' or an int device id" + ) + + def init_detector( config: str | Path | Config, checkpoint: str | None = None, palette: str = "none", - device: str = "cuda:0", + device: str | Sequence[str] | Sequence[int] = "cuda:0", cfg_options: dict | None = None, ) -> nn.Module: """Initialize a detector from config file. @@ -45,8 +76,15 @@ def init_detector( is stored in checkpoint, use checkpoint's palette first, otherwise use externally passed palette. Currently, supports 'coco', 'voc', 'citys' and 'random'. Defaults to none. - device (str): The device where the anchors will be put on. - Defaults to cuda:0. + device (str | Sequence[str] | Sequence[int]): + The device(s) to run inference on. + + - Single device examples: ``"cpu"``, ``"cuda:0"`` + - Multi-GPU single-process example: ``"cuda:0,1"`` or + ``["cuda:0", "cuda:1"]`` + + When multiple CUDA devices are provided, the model is wrapped in + :class:`visdet.engine.model.wrappers.MMDataParallel`. cfg_options (dict, optional): Options to override some settings in the used config. @@ -113,7 +151,22 @@ def init_detector( model.dataset_meta["palette"] = "random" model.cfg = config # save the config in the model for convenience - model.to(device) + + primary_device, device_ids = _parse_inference_devices(device) + model.to(primary_device) + + if device_ids is not None and len(device_ids) > 1: + if not torch.cuda.is_available(): + raise RuntimeError(f"CUDA is not available, cannot use multi-GPU device={device!r}") + if max(device_ids) >= torch.cuda.device_count(): + raise RuntimeError( + f"Requested device_ids={device_ids} but only {torch.cuda.device_count()} CUDA device(s) are available" + ) + + from visdet.engine.model import MMDataParallel + + model = MMDataParallel(model, device_ids=device_ids, output_device=device_ids[0]) + model.eval() return model @@ -162,8 +215,8 @@ def inference_detector( # RoIPool check removed - eliminating C++ ops - result_list = [] - for _i, img in enumerate(imgs): + data_list = [] + for img in imgs: # prepare data if isinstance(img, np.ndarray): # TODO: remove img_id. @@ -176,22 +229,18 @@ def inference_detector( data_["text"] = text_prompt data_["custom_entities"] = custom_entities - # build the data pipeline - data_ = test_pipeline(data_) - - data_["inputs"] = [data_["inputs"]] - data_["data_samples"] = [data_["data_samples"]] + data_list.append(test_pipeline(data_)) - # forward the model - with torch.inference_mode(): - results = model.test_step(data_)[0] + batch = pseudo_collate(data_list) - result_list.append(results) + # forward the model + with torch.inference_mode(): + results = model.test_step(batch) if not is_batch: - return result_list[0] + return results[0] else: - return result_list + return results # TODO: Awaiting refactoring diff --git a/visdet/engine/model/__init__.py b/visdet/engine/model/__init__.py index 6f7790f4..ec7a89ee 100644 --- a/visdet/engine/model/__init__.py +++ b/visdet/engine/model/__init__.py @@ -34,6 +34,7 @@ xavier_init, ) from visdet.engine.model.wrappers import ( + MMDataParallel, MMDistributedDataParallel, MMSeparateDistributedDataParallel, is_model_wrapper, @@ -49,6 +50,7 @@ # "ExponentialMovingAverage", # Not imported - EMA is in hooks module "ImgDataPreprocessor", "KaimingInit", + "MMDataParallel", "MMDistributedDataParallel", "MMSeparateDistributedDataParallel", "ModuleDict", diff --git a/visdet/engine/model/wrappers/__init__.py b/visdet/engine/model/wrappers/__init__.py index bbf728cb..1d25845d 100644 --- a/visdet/engine/model/wrappers/__init__.py +++ b/visdet/engine/model/wrappers/__init__.py @@ -3,11 +3,12 @@ # Copyright (c) OpenMMLab. All rights reserved. from visdet.engine.utils.dl_utils import TORCH_VERSION from visdet.engine.utils.version_utils import digit_version -from visdet.engine.model.wrappers.distributed import MMDistributedDataParallel +from visdet.engine.model.wrappers.distributed import MMDataParallel, MMDistributedDataParallel from visdet.engine.model.wrappers.seperate_distributed import MMSeparateDistributedDataParallel from visdet.engine.model.wrappers.utils import is_model_wrapper __all__ = [ + "MMDataParallel", "MMDistributedDataParallel", "MMSeparateDistributedDataParallel", "is_model_wrapper", diff --git a/visdet/engine/model/wrappers/distributed.py b/visdet/engine/model/wrappers/distributed.py index e4e1714e..5236b476 100644 --- a/visdet/engine/model/wrappers/distributed.py +++ b/visdet/engine/model/wrappers/distributed.py @@ -18,6 +18,29 @@ MODEL_WRAPPERS.register_module(module=DataParallel) +@MODEL_WRAPPERS.register_module(force=True) +class MMDataParallel(DataParallel): + """A data parallel model wrapper for inference. + + Compared with :class:`torch.nn.parallel.DataParallel`, this wrapper adds + :meth:`val_step` and :meth:`test_step` so it can be used with visdet's + higher-level APIs that expect these methods. + + Note: + This wrapper parallelizes the model forward pass across multiple GPUs + within a single process. For best performance, pass a batch of images + (e.g. a list of image paths/arrays) to :func:`visdet.apis.inference_detector`. + """ + + def val_step(self, data: dict | tuple | list) -> list: + data = self.module.data_preprocessor(data, False) + return self(**data, mode="predict") + + def test_step(self, data: dict | tuple | list) -> list: + data = self.module.data_preprocessor(data, False) + return self(**data, mode="predict") + + @MODEL_WRAPPERS.register_module(force=True) class MMDistributedDataParallel(DistributedDataParallel): """A distributed model wrapper used for training,testing and validation in