Fix cpu inference (open-mmlab#152)

* Add missing map_location * Add docstring * Update mmseg/apis/inference.py Co-authored-by: Jerry Jiarui XU <[email protected]> * Update inference.py * Update inference.py Co-authored-by: Jerry Jiarui XU <[email protected]>
bowenroom · Sep 24, 2020 · e385842 · e385842
1 parent 51e4cde
commit e385842
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/mmseg/apis/inference.py b/mmseg/apis/inference.py
@@ -16,7 +16,8 @@ def init_segmentor(config, checkpoint=None, device='cuda:0'):
             object.
         checkpoint (str, optional): Checkpoint path. If left as None, the model
             will not load any weights.
-
+        device (str, optional) CPU/CUDA device option. Default 'cuda:0'.
+            Use 'cpu' for loading model on CPU.
     Returns:
         nn.Module: The constructed segmentor.
     """
@@ -28,7 +29,7 @@ def init_segmentor(config, checkpoint=None, device='cuda:0'):
     config.model.pretrained = None
     model = build_segmentor(config.model, test_cfg=config.test_cfg)
     if checkpoint is not None:
-        checkpoint = load_checkpoint(model, checkpoint)
+        checkpoint = load_checkpoint(model, checkpoint, map_location='cpu')
         model.CLASSES = checkpoint['meta']['CLASSES']
         model.PALETTE = checkpoint['meta']['PALETTE']
     model.cfg = config  # save the config in the model for convenience