Support mim train / test / gridsearch in OpenMMLab-2.0 (#160)

* Suport openmmlab-2.0 with the following features:

- import Config from mmengine
- remove extra options for train / test / gridsearch
- create destination directory if dose not exists in `mim donwload`
- add mmengine in `mim list`

* fix ci with the following 2 problems:
- mmcv2.x have not pre-built package for torch1.5.
- patch xxx not found after install (cause by importlib.reload).

* Update mim/commands/download.py

Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

* dev(download): remove echo line

* support openmmlab 1.0 and 2.0

* use tests/data/lenet5_mnist_2.0.py in test case

* use tests/data/lenet5_mnist_2.0.py in test case

Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

.github/workflows/build.yml

@@ -32,10 +32,10 @@ jobs:
     strategy:
       matrix:
         python-version: [3.7]
-        torch: [1.5.0+cu101, 1.8.0+cu101]
+        torch: [1.6.0+cu101, 1.8.0+cu101]
         include:
-          - torch: 1.5.0+cu101
-            torchvision: 0.6.0+cu101
+          - torch: 1.6.0+cu101
+            torchvision: 0.7.0+cu101
           - torch: 1.8.0+cu101
             torchvision: 0.9.0+cu101
           - torch: 1.8.0+cu101
@@ -98,10 +98,10 @@ jobs:
     strategy:
       matrix:
         python-version: [3.7]
-        torch: [1.5.0, 1.8.0]
+        torch: [1.6.0, 1.8.0]
         include:
-          - torch: 1.5.0
-            torchvision: 0.6.0
+          - torch: 1.6.0
+            torchvision: 0.7.0
           - torch: 1.8.0
             torchvision: 0.9.0
     steps:

mim/commands/download.py

@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import os
 import os.path as osp
 from typing import List, Optional
 
@@ -65,6 +66,10 @@ def download(package: str,
 
     dest_root = osp.abspath(dest_root)
 
+    # Create the destination directory if it does not exist.
+    if not osp.exists(dest_root):
+        os.makedirs(dest_root)
+
     package, version = split_package_version(package)
     if version:
         raise ValueError(
@@ -89,8 +94,11 @@ def download(package: str,
     try:
         from mmengine import Config
     except ImportError:
-        msg = 'Please install mmengine to use the download command.'
-        raise ImportError(highlighted_error(msg))
+        try:
+            from mmcv import Config
+        except ImportError:
+            raise ImportError(
+                'Please install mmengine to use the download command!')
 
     for config in configs:
         click.echo(f'processing {config}...')

mim/commands/gridsearch.py

@@ -272,8 +272,11 @@ def gridsearch(
     try:
         from mmengine import Config
     except ImportError:
-        msg = 'Please install mmengine to use the gridsearch command.'
-        raise ImportError(highlighted_error(msg))
+        try:
+            from mmcv import Config
+        except ImportError:
+            raise ImportError(
+                'Please install mmengine to use the gridsearch command!')
 
     cfg = Config.fromfile(config)
     for arg in search_args_dict:
@@ -354,14 +357,14 @@ def gridsearch(
         common_args = ['--launcher', launcher] + other_args_str.split()
 
         if launcher == 'none':
-            if gpus:
-                cmd = [
-                    'python', train_script, config_path, '--gpus',
-                    str(gpus)
-                ] + common_args
-            else:
-                cmd = ['python', train_script, config_path, '--device', 'cpu'
-                       ] + common_args
+            cmd = ['python', train_script, config_path] + common_args
+            help_msg = subprocess.check_output(['python', train_script, '-h'])
+            if '--gpus' in help_msg.decode():
+                # OpenMMLab 1.0 should add the `--gpus` or `--device` flags.
+                if gpus:
+                    cmd += ['--gpus', str(gpus)]
+                else:
+                    cmd += ['--device', 'cpu']
         elif launcher == 'pytorch':
             cmd = [
                 'python', '-m', 'torch.distributed.launch',

mim/commands/install.py

@@ -89,10 +89,22 @@ def install(
     """
 
     # Reload `pip._vendor.pkg_resources` so that pip can refresh to get the
-    # latest working set.
+    # latest working set in the same process.
     # In some cases, when a package is uninstalled and then installed, the
     # working set is not updated in time, leading to the mistaken belief that
     # the package is already installed.
+    #
+    # NOTE: Some unpredictable bugs could occurs with `importlib.reload`.
+    # A known issues in pip < 22.0: `METADATA not found in /tmp/xxx/xxx.whel`
+    # >>> import pip._vendor.pkg_resources
+    # >>> import importlib
+    # >>> a = pip._vendor.pkg_resources.DistInfoDistribution()
+    # >>> type(a) is pip._vendor.pkg_resources.DistInfoDistribution
+    # True
+    # >>> importlib.reload(pip._vendor.pkg_resources)
+    # <module 'pip._vendor.pkg_resources' from '...'>
+    # >>> type(a) is pip._vendor.pkg_resources.DistInfoDistribution
+    # False  # This will cause some problems!!!
     importlib.reload(pip._vendor.pkg_resources)
 
     # Get mmcv_base_url from environment variable if exists.

mim/commands/list.py

@@ -66,7 +66,7 @@ def list_package(all: bool = False) -> List[Tuple[str, ...]]:
                 feed_parser.feed(metadata)
                 home_page = feed_parser.close().get('home-page')
 
-            if pkg_name.startswith('mmcv'):
+            if pkg_name.startswith('mmcv') or pkg_name == 'mmengine':
                 pkgs_info.append((pkg_name, pkg.version, home_page))
                 continue
 

mim/commands/train.py

@@ -225,12 +225,14 @@ def train(
     common_args = ['--launcher', launcher] + list(other_args)
 
     if launcher == 'none':
-        if gpus:
-            cmd = ['python', train_script, config, '--gpus',
-                   str(gpus)] + common_args
-        else:
-            cmd = ['python', train_script, config, '--device', 'cpu'
-                   ] + common_args
+        cmd = ['python', train_script, config] + common_args
+        help_msg = subprocess.check_output(['python', train_script, '-h'])
+        if '--gpus' in help_msg.decode():
+            # OpenMMLab 1.0 should add the `--gpus` or `--device` flags.
+            if gpus:
+                cmd += ['--gpus', str(gpus)]
+            else:
+                cmd += ['--device', 'cpu']
     elif launcher == 'pytorch':
         cmd = [
             'python', '-m', 'torch.distributed.launch',

tests/data/lenet5_mnist_2.0.py

@@ -0,0 +1,91 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+# model settings
+model = dict(
+    type='ImageClassifier',
+    backbone=dict(type='LeNet5', num_classes=10),
+    neck=None,
+    head=dict(
+        type='ClsHead',
+        loss=dict(type='CrossEntropyLoss', loss_weight=1.0),
+    ))
+# dataset settings
+dataset_type = 'MNIST'
+data_preprocessor = dict(mean=[33.46], std=[78.87])
+
+pipeline = [dict(type='Resize', scale=32), dict(type='PackClsInputs')]
+
+common_data_cfg = dict(
+    type=dataset_type, data_prefix='data/mnist', pipeline=pipeline)
+
+train_dataloader = dict(
+    batch_size=128,
+    num_workers=2,
+    dataset=dict(**common_data_cfg, test_mode=False),
+    sampler=dict(type='DefaultSampler', shuffle=True),
+    persistent_workers=True,
+)
+
+val_dataloader = dict(
+    batch_size=128,
+    num_workers=2,
+    dataset=dict(**common_data_cfg, test_mode=True),
+    sampler=dict(type='DefaultSampler', shuffle=False),
+    persistent_workers=True,
+)
+val_evaluator = dict(type='Accuracy', topk=(1, ))
+
+test_dataloader = val_dataloader
+test_evaluator = val_evaluator
+
+# schedule settings
+optim_wrapper = dict(
+    optimizer=dict(type='SGD', lr=0.01, momentum=0.9, weight_decay=0.0001))
+
+param_scheduler = dict(
+    type='MultiStepLR',  # learning policy, decay on several milestones.
+    by_epoch=True,  # update based on epoch.
+    milestones=[15],  # decay at the 15th epochs.
+    gamma=0.1,  # decay to 0.1 times.
+)
+
+train_cfg = dict(by_epoch=True, max_epochs=1, val_interval=2)  # train 5 epochs
+val_cfg = dict()
+test_cfg = dict()
+
+# runtime settings
+default_scope = 'mmcls'
+
+default_hooks = dict(
+    # record the time of every iteration.
+    timer=dict(type='IterTimerHook'),
+    # print log every 150 iterations.
+    logger=dict(type='LoggerHook', interval=150),
+    # enable the parameter scheduler.
+    param_scheduler=dict(type='ParamSchedulerHook'),
+    # save checkpoint per epoch.
+    checkpoint=dict(type='CheckpointHook', interval=1),
+    # set sampler seed in distributed evrionment.
+    sampler_seed=dict(type='DistSamplerSeedHook'),
+)
+
+env_cfg = dict(
+    # disable cudnn benchmark
+    cudnn_benchmark=False,
+    # set multi process parameters
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    # set distributed parameters
+    dist_cfg=dict(backend='nccl'),
+)
+
+log_level = 'INFO'
+
+# load from which checkpoint
+load_from = None
+
+# whether to resume the training of the checkpoint
+resume_from = None
+
+# NOTE: `auto_scale_lr` is for automatically scaling LR
+# based on the actual training batch size.
+# base_batch_size = (1 GPUs) x (128 samples per GPU)
+auto_scale_lr = dict(base_batch_size=128)

tests/test_gridsearch.py

@@ -13,6 +13,8 @@ def setup_module():
     runner = CliRunner()
     result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
     assert result.exit_code == 0
+    result = runner.invoke(uninstall, ['mmcv', '--yes'])
+    assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmengine', '--yes'])
     assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmcls', '--yes'])
@@ -28,30 +30,30 @@ def setup_module():
 ])
 def test_gridsearch(gpus, tmp_path):
     runner = CliRunner()
-    result = runner.invoke(install, ['mmcls', '--yes'])
+    result = runner.invoke(install, ['mmcls>=1.0.0rc0', '--yes'])
     assert result.exit_code == 0
     result = runner.invoke(install, ['mmengine', '--yes'])
     assert result.exit_code == 0
-    # Since `mminstall.txt` is not included in the distribution of
-    # mmcls<=0.23.1, we need to install mmcv-full manually.
-    result = runner.invoke(install, ['mmcv-full', '--yes'])
+    result = runner.invoke(install, ['mmcv>=2.0.0rc0', '--yes'])
     assert result.exit_code == 0
 
     args1 = [
-        'mmcls', 'tests/data/lenet5_mnist.py', f'--gpus={gpus}',
-        f'--work-dir={tmp_path}', '--search-args', '--optimizer.lr 1e-3 1e-4'
+        'mmcls', 'tests/data/lenet5_mnist_2.0.py', f'--gpus={gpus}',
+        f'--work-dir={tmp_path}', '--search-args',
+        '--optim_wrapper.optimizer.lr 1e-3 1e-4'
     ]
     args2 = [
-        'mmcls', 'tests/data/lenet5_mnist.py', f'--gpus={gpus}',
+        'mmcls', 'tests/data/lenet5_mnist_2.0.py', f'--gpus={gpus}',
         f'--work-dir={tmp_path}', '--search-args',
-        '--optimizer.weight_decay 1e-3 1e-4'
+        '--optim_wrapper.optimizer.weight_decay 1e-3 1e-4'
     ]
     args3 = [
         'mmcls', 'tests/data/xxx.py', f'--gpus={gpus}',
-        f'--work-dir={tmp_path}', '--search-args', '--optimizer.lr 1e-3 1e-4'
+        f'--work-dir={tmp_path}', '--search-args',
+        '--optim_wrapper.optimizer.lr 1e-3 1e-4'
     ]
     args4 = [
-        'mmcls', 'tests/data/lenet5_mnist.py', f'--gpus={gpus}',
+        'mmcls', 'tests/data/lenet5_mnist_2.0.py', f'--gpus={gpus}',
         f'--work-dir={tmp_path}', '--search-args'
     ]
 
@@ -72,6 +74,8 @@ def teardown_module():
     runner = CliRunner()
     result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
     assert result.exit_code == 0
+    result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
+    assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmengine', '--yes'])
     assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmcls', '--yes'])

tests/test_list.py

@@ -10,6 +10,8 @@ def setup_module():
     runner = CliRunner()
     result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
     assert result.exit_code == 0
+    result = runner.invoke(uninstall, ['mmcv', '--yes'])
+    assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmcls', '--yes'])
     assert result.exit_code == 0
 
@@ -30,5 +32,7 @@ def teardown_module():
     runner = CliRunner()
     result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
     assert result.exit_code == 0
+    result = runner.invoke(uninstall, ['mmcv', '--yes'])
+    assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmcls', '--yes'])
     assert result.exit_code == 0

tests/test_run.py

@@ -12,6 +12,8 @@ def setup_module():
     runner = CliRunner()
     result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
     assert result.exit_code == 0
+    result = runner.invoke(uninstall, ['mmcv', '--yes'])
+    assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmcls', '--yes'])
     assert result.exit_code == 0
 
@@ -26,32 +28,30 @@ def setup_module():
 ])
 def test_run(device, gpus, tmp_path):
     runner = CliRunner()
-    result = runner.invoke(install, ['mmcls', '--yes'])
+    result = runner.invoke(install, ['mmcls>=1.0.0rc0', '--yes'])
     assert result.exit_code == 0
-    # Since `mminstall.txt` is not included in the distribution of
-    # mmcls<=0.23.1, we need to install mmcv-full manually.
-    result = runner.invoke(install, ['mmcv-full', '--yes'])
+    result = runner.invoke(install, ['mmengine', '--yes'])
+    assert result.exit_code == 0
+    result = runner.invoke(install, ['mmcv>=2.0.0rc0', '--yes'])
     assert result.exit_code == 0
 
     result = runner.invoke(run, [
-        'mmcls', 'train', 'tests/data/lenet5_mnist.py', f'--gpus={gpus}',
+        'mmcls', 'train', 'tests/data/lenet5_mnist_2.0.py',
         f'--work-dir={tmp_path}'
     ])
     assert result.exit_code == 0
     result = runner.invoke(run, [
-        'mmcls', 'test', 'tests/data/lenet5_mnist.py',
-        'tests/data/epoch_1.pth', f'--device={device}', '--metrics=accuracy'
+        'mmcls', 'test', 'tests/data/lenet5_mnist_2.0.py',
+        'tests/data/epoch_1.pth'
     ])
     assert result.exit_code == 0
     result = runner.invoke(run, [
-        'mmcls', 'xxx', 'tests/data/lenet5_mnist.py', 'tests/data/epoch_1.pth',
-        f'--gpus={gpus}', '--metrics=accuracy'
+        'mmcls', 'xxx', 'tests/data/lenet5_mnist_2.0.py',
+        'tests/data/epoch_1.pth'
     ])
     assert result.exit_code != 0
-    result = runner.invoke(run, [
-        'mmcls', 'test', 'tests/data/xxx.py', 'tests/data/epoch_1.pth',
-        f'--device={device}', '--metrics=accuracy'
-    ])
+    result = runner.invoke(
+        run, ['mmcls', 'test', 'tests/data/xxx.py', 'tests/data/epoch_1.pth'])
     assert result.exit_code != 0
 
 
@@ -59,5 +59,7 @@ def teardown_module():
     runner = CliRunner()
     result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
     assert result.exit_code == 0
+    result = runner.invoke(uninstall, ['mmcv', '--yes'])
+    assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmcls', '--yes'])
     assert result.exit_code == 0

tests/test_search.py

@@ -13,6 +13,8 @@ def setup_module():
     runner = CliRunner()
     result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
     assert result.exit_code == 0
+    result = runner.invoke(uninstall, ['mmcv', '--yes'])
+    assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmcls', '--yes'])
     assert result.exit_code == 0
 
@@ -119,5 +121,7 @@ def teardown_module():
     runner = CliRunner()
     result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
     assert result.exit_code == 0
+    result = runner.invoke(uninstall, ['mmcv', '--yes'])
+    assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmcls', '--yes'])
     assert result.exit_code == 0

tests/test_test.py

@@ -12,6 +12,8 @@ def setup_module():
     runner = CliRunner()
     result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
     assert result.exit_code == 0
+    result = runner.invoke(uninstall, ['mmcv', '--yes'])
+    assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmcls', '--yes'])
     assert result.exit_code == 0
 
@@ -25,26 +27,25 @@ def setup_module():
 ])
 def test_test(device):
     runner = CliRunner()
-    result = runner.invoke(install, ['mmcls', '--yes'])
+    result = runner.invoke(install, ['mmcls>=1.0.0rc0', '--yes'])
     assert result.exit_code == 0
-    # Since `mminstall.txt` is not included in the distribution of
-    # mmcls<=0.23.1, we need to install mmcv-full manually.
-    result = runner.invoke(install, ['mmcv-full', '--yes'])
+    result = runner.invoke(install, ['mmengine', '--yes'])
+    assert result.exit_code == 0
+    result = runner.invoke(install, ['mmcv>=2.0.0rc0', '--yes'])
     assert result.exit_code == 0
 
     result = runner.invoke(test, [
-        'mmcls', 'tests/data/lenet5_mnist.py', '--checkpoint',
-        'tests/data/epoch_1.pth', f'--device={device}', '--metrics=accuracy'
+        'mmcls', 'tests/data/lenet5_mnist_2.0.py', '--checkpoint',
+        'tests/data/epoch_1.pth'
     ])
     assert result.exit_code == 0
     result = runner.invoke(test, [
-        'mmcls', 'tests/data/xxx.py', '--checkpoint', 'tests/data/epoch_1.pth',
-        f'--device={device}', '--metrics=accuracy'
+        'mmcls', 'tests/data/xxx.py', '--checkpoint', 'tests/data/epoch_1.pth'
     ])
     assert result.exit_code != 0
     result = runner.invoke(test, [
-        'mmcls', 'tests/data/lenet5_mnist.py', '--checkpoint',
-        'tests/data/xxx.pth', f'--device={device}', '--metrics=accuracy'
+        'mmcls', 'tests/data/lenet5_mnist_2.0.py', '--checkpoint',
+        'tests/data/xxx.pth'
     ])
     assert result.exit_code != 0
 
@@ -53,5 +54,7 @@ def teardown_module():
     runner = CliRunner()
     result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
     assert result.exit_code == 0
+    result = runner.invoke(uninstall, ['mmcv', '--yes'])
+    assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmcls', '--yes'])
     assert result.exit_code == 0

tests/test_train.py

@@ -12,6 +12,8 @@ def setup_module():
     runner = CliRunner()
     result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
     assert result.exit_code == 0
+    result = runner.invoke(uninstall, ['mmcv', '--yes'])
+    assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmcls', '--yes'])
     assert result.exit_code == 0
 
@@ -25,15 +27,15 @@ def setup_module():
 ])
 def test_train(gpus, tmp_path):
     runner = CliRunner()
-    result = runner.invoke(install, ['mmcls', '--yes'])
+    result = runner.invoke(install, ['mmcls>=1.0.0rc0', '--yes'])
     assert result.exit_code == 0
-    # Since `mminstall.txt` is not included in the distribution of
-    # mmcls<=0.23.1, we need to install mmcv-full manually.
-    result = runner.invoke(install, ['mmcv-full', '--yes'])
+    result = runner.invoke(install, ['mmengine', '--yes'])
+    assert result.exit_code == 0
+    result = runner.invoke(install, ['mmcv>=2.0.0rc0', '--yes'])
     assert result.exit_code == 0
 
     result = runner.invoke(train, [
-        'mmcls', 'tests/data/lenet5_mnist.py', f'--gpus={gpus}',
+        'mmcls', 'tests/data/lenet5_mnist_2.0.py', f'--gpus={gpus}',
         f'--work-dir={tmp_path}'
     ])
     assert result.exit_code == 0
@@ -49,5 +51,7 @@ def teardown_module():
     runner = CliRunner()
     result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
     assert result.exit_code == 0
+    result = runner.invoke(uninstall, ['mmcv', '--yes'])
+    assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmcls', '--yes'])
     assert result.exit_code == 0

tests/test_utils.py

@@ -10,6 +10,8 @@ def setup_module():
     runner = CliRunner()
     result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
     assert result.exit_code == 0
+    result = runner.invoke(uninstall, ['mmcv', '--yes'])
+    assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmcls', '--yes'])
     assert result.exit_code == 0
 
@@ -41,5 +43,7 @@ def teardown_module():
     runner = CliRunner()
     result = runner.invoke(uninstall, ['mmcv-full', '--yes'])
     assert result.exit_code == 0
+    result = runner.invoke(uninstall, ['mmcv', '--yes'])
+    assert result.exit_code == 0
     result = runner.invoke(uninstall, ['mmcls', '--yes'])
     assert result.exit_code == 0