From 0f90a0af0a9fe863b5b910dc658d376f9c50952c Mon Sep 17 00:00:00 2001
From: hanrui1sensetime <83800577+hanrui1sensetime@users.noreply.github.com>
Date: Fri, 10 Dec 2021 14:15:08 +0800
Subject: [PATCH] [Fix] fix bugs for mmcls performance test (#269)

* fix bugs for mmcls performance test

* fix yapf

* add comments of CLASSES attribute
---
 mmdeploy/codebase/mmcls/deploy/classification_model.py | 5 ++++-
 mmdeploy/codebase/mmcls/deploy/mmclassification.py     | 5 ++++-
 tools/test.py                                          | 6 ++++++
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/mmdeploy/codebase/mmcls/deploy/classification_model.py b/mmdeploy/codebase/mmcls/deploy/classification_model.py
index 0c51b4635c..ed5405f786 100644
--- a/mmdeploy/codebase/mmcls/deploy/classification_model.py
+++ b/mmdeploy/codebase/mmcls/deploy/classification_model.py
@@ -61,7 +61,10 @@ def forward(self, img: List[torch.Tensor], *args, **kwargs) -> list:
             list: A list contains predictions.
         """
 
-        input_img = img[0].contiguous()
+        if isinstance(img, list):
+            input_img = img[0].contiguous()
+        else:
+            input_img = img.contiguous()
         outputs = self.forward_test(input_img, *args, **kwargs)
 
         return list(outputs)
diff --git a/mmdeploy/codebase/mmcls/deploy/mmclassification.py b/mmdeploy/codebase/mmcls/deploy/mmclassification.py
index aa683e4d89..4a637b8448 100644
--- a/mmdeploy/codebase/mmcls/deploy/mmclassification.py
+++ b/mmdeploy/codebase/mmcls/deploy/mmclassification.py
@@ -123,6 +123,7 @@ def single_gpu_test(model: torch.nn.Module,
                         data_loader: DataLoader,
                         show: bool = False,
                         out_dir: Optional[str] = None,
+                        win_name: str = '',
                         **kwargs) -> List:
         """Run test with single gpu.
 
@@ -132,10 +133,12 @@ def single_gpu_test(model: torch.nn.Module,
             show (bool): Specifying whether to show plotted results.
                 Default: False.
             out_dir (str): A directory to save results, Default: None.
+            win_name (str): The name of windows, Default: ''.
 
         Returns:
             list: The prediction results.
         """
         from mmcls.apis import single_gpu_test
-        outputs = single_gpu_test(model, data_loader, show, out_dir, **kwargs)
+        outputs = single_gpu_test(
+            model, data_loader, show, out_dir, win_name=win_name, **kwargs)
         return outputs
diff --git a/tools/test.py b/tools/test.py
index 356825cf9a..2cd4dd21ec 100644
--- a/tools/test.py
+++ b/tools/test.py
@@ -107,6 +107,12 @@ def main():
     device_id = parse_device_id(args.device)
 
     model = MMDataParallel(model, device_ids=[0])
+    # The whole dataset test wrapped a MMDataParallel class outside the module.
+    # As mmcls.apis.test.py single_gpu_test defined, the MMDataParallel needs
+    # a 'CLASSES' attribute. So we ensure the MMDataParallel class has the same
+    # CLASSES attribute as the inside module.
+    if hasattr(model.module, 'CLASSES'):
+        model.CLASSES = model.module.CLASSES
     if args.speed_test:
         with_sync = device_id == 0
         output_file = sys.stdout