From 8c0e093b31b35890ebf06e61c54838134869f234 Mon Sep 17 00:00:00 2001
From: Jerry Jiarui XU <xvjiarui0826@gmail.com>
Date: Tue, 25 Aug 2020 20:01:01 +0800
Subject: [PATCH] Fixed slide inference (#90)

---
 docs/getting_started.md                    | 4 ++--
 mmseg/apis/test.py                         | 2 +-
 mmseg/datasets/cityscapes.py               | 2 +-
 mmseg/models/segmentors/encoder_decoder.py | 6 ++++--
 tools/train.py                             | 2 +-
 5 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/docs/getting_started.md b/docs/getting_started.md
index a4ab035245..892060d00b 100644
--- a/docs/getting_started.md
+++ b/docs/getting_started.md
@@ -229,7 +229,7 @@ To trade speed with GPU memory, you may pass in `--options model.backbone.with_c
 python tools/train.py ${CONFIG_FILE} [optional arguments]
 ```
 
-If you want to specify the working directory in the command, you can add an argument `--work_dir ${YOUR_WORK_DIR}`.
+If you want to specify the working directory in the command, you can add an argument `--work-dir ${YOUR_WORK_DIR}`.
 
 ### Train with multiple GPUs
 
@@ -253,7 +253,7 @@ Difference between `resume-from` and `load-from`:
 If you run MMSegmentation on a cluster managed with [slurm](https://slurm.schedmd.com/), you can use the script `slurm_train.sh`. (This script also supports single machine training.)
 
 ```shell
-[GPUS=${GPUS}] ./tools/slurm_train.sh ${PARTITION} ${JOB_NAME} ${CONFIG_FILE} ${WORK_DIR}
+[GPUS=${GPUS}] ./tools/slurm_train.sh ${PARTITION} ${JOB_NAME} ${CONFIG_FILE} --work-dir ${WORK_DIR}
 ```
 
 Here is an example of using 16 GPUs to train PSPNet on the dev partition.
diff --git a/mmseg/apis/test.py b/mmseg/apis/test.py
index 8cbf236f05..5b4b2da40e 100644
--- a/mmseg/apis/test.py
+++ b/mmseg/apis/test.py
@@ -30,7 +30,7 @@ def single_gpu_test(model, data_loader, show=False, out_dir=None):
     prog_bar = mmcv.ProgressBar(len(dataset))
     for i, data in enumerate(data_loader):
         with torch.no_grad():
-            result = model(return_loss=False, rescale=not show, **data)
+            result = model(return_loss=False, **data)
         if isinstance(results, list):
             results.extend(result)
         else:
diff --git a/mmseg/datasets/cityscapes.py b/mmseg/datasets/cityscapes.py
index 30e3c2b24e..e26cd00b09 100644
--- a/mmseg/datasets/cityscapes.py
+++ b/mmseg/datasets/cityscapes.py
@@ -173,7 +173,7 @@ def _evaluate_cityscapes(self, results, logger, imgfile_prefix):
         try:
             import cityscapesscripts.evaluation.evalPixelLevelSemanticLabeling as CSEval  # noqa
         except ImportError:
-            raise ImportError('Please run "pip install citscapesscripts" to '
+            raise ImportError('Please run "pip install cityscapesscripts" to '
                               'install cityscapesscripts first.')
         msg = 'Evaluating in Cityscapes style'
         if logger is None:
diff --git a/mmseg/models/segmentors/encoder_decoder.py b/mmseg/models/segmentors/encoder_decoder.py
index d1709e0ca3..3e11630e25 100644
--- a/mmseg/models/segmentors/encoder_decoder.py
+++ b/mmseg/models/segmentors/encoder_decoder.py
@@ -195,8 +195,10 @@ def slide_inference(self, img, img_meta, rescale):
 
                 count_mat[:, :, y1:y2, x1:x2] += 1
         assert (count_mat == 0).sum() == 0
-        # We want to regard count_mat as a constant while exporting to ONNX
-        count_mat = torch.from_numpy(count_mat.detach().numpy())
+        if torch.onnx.is_in_onnx_export():
+            # cast count_mat to constant while exporting to ONNX
+            count_mat = torch.from_numpy(
+                count_mat.cpu().detach().numpy()).to(device=img.device)
         preds = preds / count_mat
         if rescale:
             preds = resize(
diff --git a/tools/train.py b/tools/train.py
index 7f67c72004..8e3835ae00 100644
--- a/tools/train.py
+++ b/tools/train.py
@@ -19,7 +19,7 @@
 def parse_args():
     parser = argparse.ArgumentParser(description='Train a segmentor')
     parser.add_argument('config', help='train config file path')
-    parser.add_argument('--work_dir', help='the dir to save logs and models')
+    parser.add_argument('--work-dir', help='the dir to save logs and models')
     parser.add_argument(
         '--load-from', help='the checkpoint file to load weights from')
     parser.add_argument(