From dfae731e900f60e0b7ce1c40f48107fc16e8d183 Mon Sep 17 00:00:00 2001 From: David de la Iglesia Castro Date: Thu, 24 Sep 2020 08:51:03 +0200 Subject: [PATCH 1/5] Add missing map_location --- mmseg/apis/inference.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mmseg/apis/inference.py b/mmseg/apis/inference.py index 3ba6b62ce1..250f3a4cc0 100644 --- a/mmseg/apis/inference.py +++ b/mmseg/apis/inference.py @@ -28,7 +28,8 @@ 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) + map_loc = 'cpu' if device == 'cpu' else None + checkpoint = load_checkpoint(model, checkpoint, map_location=map_loc) model.CLASSES = checkpoint['meta']['CLASSES'] model.PALETTE = checkpoint['meta']['PALETTE'] model.cfg = config # save the config in the model for convenience From 9a8de03aab8a49e268273b6de88d571214d7a4d1 Mon Sep 17 00:00:00 2001 From: David de la Iglesia Castro Date: Thu, 24 Sep 2020 08:53:02 +0200 Subject: [PATCH 2/5] Add docstring --- mmseg/apis/inference.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/mmseg/apis/inference.py b/mmseg/apis/inference.py index 250f3a4cc0..7653cd19b8 100644 --- 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. """ From 5c7185a071e2de3375ee48cf8fe71097592e67f1 Mon Sep 17 00:00:00 2001 From: David de la Iglesia Castro Date: Thu, 24 Sep 2020 18:56:52 +0200 Subject: [PATCH 3/5] Update mmseg/apis/inference.py Co-authored-by: Jerry Jiarui XU --- mmseg/apis/inference.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mmseg/apis/inference.py b/mmseg/apis/inference.py index 7653cd19b8..11eb3f3f9a 100644 --- a/mmseg/apis/inference.py +++ b/mmseg/apis/inference.py @@ -30,7 +30,7 @@ def init_segmentor(config, checkpoint=None, device='cuda:0'): model = build_segmentor(config.model, test_cfg=config.test_cfg) if checkpoint is not None: map_loc = 'cpu' if device == 'cpu' else None - checkpoint = load_checkpoint(model, checkpoint, map_location=map_loc) + 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 From 84a42845fc5cd665904e3fcd3687d8d09f61c7a6 Mon Sep 17 00:00:00 2001 From: David de la Iglesia Castro Date: Thu, 24 Sep 2020 18:57:57 +0200 Subject: [PATCH 4/5] Update inference.py --- mmseg/apis/inference.py | 1 - 1 file changed, 1 deletion(-) diff --git a/mmseg/apis/inference.py b/mmseg/apis/inference.py index 11eb3f3f9a..6fa7e3b343 100644 --- a/mmseg/apis/inference.py +++ b/mmseg/apis/inference.py @@ -29,7 +29,6 @@ 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: - map_loc = 'cpu' if device == 'cpu' else None checkpoint = load_checkpoint(model, checkpoint, map_location='cpu') model.CLASSES = checkpoint['meta']['CLASSES'] model.PALETTE = checkpoint['meta']['PALETTE'] From cdd346a59a8bfac104a1424c6a39dbe90c74ae0c Mon Sep 17 00:00:00 2001 From: David de la Iglesia Castro Date: Thu, 24 Sep 2020 18:58:08 +0200 Subject: [PATCH 5/5] Update inference.py