mirrors/mmsegmentation
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmsegmentation.git synced 2025-06-03 22:03:48 +08:00
Code Issues Projects Releases Wiki Activity

Fix cpu inference (#152)

Browse Source 
* Add missing map_location

* Add docstring

* Update mmseg/apis/inference.py

Co-authored-by: Jerry Jiarui XU <xvjiarui0826@gmail.com>

* Update inference.py

* Update inference.py

Co-authored-by: Jerry Jiarui XU <xvjiarui0826@gmail.com>
This commit is contained in:
David de la Iglesia Castro 2020-09-24 19:34:40 +02:00 committed by GitHub
parent 746c8d3785
commit 7f8bc7935c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mmseg/apis/inference.py
View File

@ -16,7 +16,8 @@ def init_segmentor(config, checkpoint=None, device='cuda:0'):
object. object.
checkpoint (str, optional): Checkpoint path. If left as None, the model checkpoint (str, optional): Checkpoint path. If left as None, the model
will not load any weights. will not load any weights.
device (str, optional) CPU/CUDA device option. Default 'cuda:0'.
Use 'cpu' for loading model on CPU.
Returns: Returns:
nn.Module: The constructed segmentor. nn.Module: The constructed segmentor.
""" """
@ -28,7 +29,7 @@ def init_segmentor(config, checkpoint=None, device='cuda:0'):
config.model.pretrained = None config.model.pretrained = None
model = build_segmentor(config.model, test_cfg=config.test_cfg) model = build_segmentor(config.model, test_cfg=config.test_cfg)
if checkpoint is not None: if checkpoint is not None:
checkpoint = load_checkpoint(model, checkpoint) checkpoint = load_checkpoint(model, checkpoint, map_location='cpu')
model.CLASSES = checkpoint['meta']['CLASSES'] model.CLASSES = checkpoint['meta']['CLASSES']
model.PALETTE = checkpoint['meta']['PALETTE'] model.PALETTE = checkpoint['meta']['PALETTE']
model.cfg = config # save the config in the model for convenience model.cfg = config # save the config in the model for convenience