Upload 159 files

mmpl/datasets/__init__.py

@@ -0,0 +1,9 @@
+from .builder import build_dataset
+from .pl_datamodule import PLDataModule
+from .nwpu_ins_dataset import NWPUInsSegDataset
+from .whu_ins_dataset import WHUInsSegDataset
+from .ssdd_ins_dataset import SSDDInsSegDataset
+
+__all__ = [
+    'build_dataset', 'PLDataModule',
+]

mmpl/datasets/base_dataset.py

@@ -0,0 +1,212 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import os.path as osp
+from os import PathLike
+from typing import List, Optional, Sequence, Union
+
+import mmengine
+import numpy as np
+from mmengine.dataset import BaseDataset as _BaseDataset
+
+from .builder import DATASETS
+
+
+def expanduser(path):
+    """Expand ~ and ~user constructions.
+
+    If user or $HOME is unknown, do nothing.
+    """
+    if isinstance(path, (str, PathLike)):
+        return osp.expanduser(path)
+    else:
+        return path
+
+
+@DATASETS.register_module()
+class BaseDataset(_BaseDataset):
+    """Base dataset for image classification task.
+
+    This dataset support annotation file in `OpenMMLab 2.0 style annotation
+    format`.
+
+    .. _OpenMMLab 2.0 style annotation format:
+        https://github.com/open-mmlab/mmengine/blob/main/docs/zh_cn/tutorials/basedataset.md
+
+    Comparing with the :class:`mmengine.BaseDataset`, this class implemented
+    several useful methods.
+
+    Args:
+        ann_file (str): Annotation file path.
+        metainfo (dict, optional): Meta information for dataset, such as class
+            information. Defaults to None.
+        data_root (str): The root directory for ``data_prefix`` and
+            ``ann_file``. Defaults to ''.
+        data_prefix (str | dict): Prefix for training data. Defaults to ''.
+        filter_cfg (dict, optional): Config for filter data. Defaults to None.
+        indices (int or Sequence[int], optional): Support using first few
+            data in annotation file to facilitate training/testing on a smaller
+            dataset. Defaults to None, which means using all ``data_infos``.
+        serialize_data (bool): Whether to hold memory using serialized objects,
+            when enabled, data loader workers can use shared RAM from master
+            process instead of making a copy. Defaults to True.
+        pipeline (Sequence): Processing pipeline. Defaults to an empty tuple.
+        test_mode (bool): ``test_mode=True`` means in test phase.
+            Defaults to False.
+        lazy_init (bool): Whether to load annotation during instantiation.
+            In some cases, such as visualization, only the meta information of
+            the dataset is needed, which is not necessary to load annotation
+            file. ``Basedataset`` can skip load annotations to save time by set
+            ``lazy_init=False``. Defaults to False.
+        max_refetch (int): If ``Basedataset.prepare_data`` get a None img.
+            The maximum extra number of cycles to get a valid image.
+            Defaults to 1000.
+        classes (str | Sequence[str], optional): Specify names of classes.
+
+            - If is string, it should be a file path, and the every line of
+              the file is a name of a class.
+            - If is a sequence of string, every item is a name of class.
+            - If is None, use categories information in ``metainfo`` argument,
+              annotation file or the class attribute ``METAINFO``.
+
+            Defaults to None.
+    """  # noqa: E501
+
+    def __init__(self,
+                 ann_file: str = '',
+                 metainfo: Optional[dict] = None,
+                 data_root: str = '',
+                 data_prefix: Union[str, dict] = '',
+                 filter_cfg: Optional[dict] = None,
+                 indices: Optional[Union[int, Sequence[int]]] = None,
+                 serialize_data: bool = True,
+                 pipeline: Sequence = (),
+                 test_mode: bool = False,
+                 lazy_init: bool = False,
+                 max_refetch: int = 1000,
+                 classes: Union[str, Sequence[str], None] = None):
+        if isinstance(data_prefix, str):
+            data_prefix = dict(img_path=expanduser(data_prefix))
+
+        ann_file = expanduser(ann_file)
+        metainfo = self._compat_classes(metainfo, classes)
+
+        super().__init__(
+            ann_file=ann_file,
+            metainfo=metainfo,
+            data_root=data_root,
+            data_prefix=data_prefix,
+            filter_cfg=filter_cfg,
+            indices=indices,
+            serialize_data=serialize_data,
+            pipeline=pipeline,
+            test_mode=test_mode,
+            lazy_init=lazy_init,
+            max_refetch=max_refetch)
+
+    @property
+    def img_prefix(self):
+        """The prefix of images."""
+        return self.data_prefix['img_path']
+
+    @property
+    def CLASSES(self):
+        """Return all categories names."""
+        return self._metainfo.get('classes', None)
+
+    @property
+    def class_to_idx(self):
+        """Map mapping class name to class index.
+
+        Returns:
+            dict: mapping from class name to class index.
+        """
+
+        return {cat: i for i, cat in enumerate(self.CLASSES)}
+
+    def get_gt_labels(self):
+        """Get all ground-truth labels (categories).
+
+        Returns:
+            np.ndarray: categories for all images.
+        """
+
+        gt_labels = np.array(
+            [self.get_data_info(i)['gt_label'] for i in range(len(self))])
+        return gt_labels
+
+    def get_cat_ids(self, idx: int) -> List[int]:
+        """Get category id by index.
+
+        Args:
+            idx (int): Index of data.
+
+        Returns:
+            cat_ids (List[int]): Image category of specified index.
+        """
+
+        return [int(self.get_data_info(idx)['gt_label'])]
+
+    def _compat_classes(self, metainfo, classes):
+        """Merge the old style ``classes`` arguments to ``metainfo``."""
+        if isinstance(classes, str):
+            # take it as a file path
+            class_names = mmengine.list_from_file(expanduser(classes))
+        elif isinstance(classes, (tuple, list)):
+            class_names = classes
+        elif classes is not None:
+            raise ValueError(f'Unsupported type {type(classes)} of classes.')
+
+        if metainfo is None:
+            metainfo = {}
+
+        if classes is not None:
+            metainfo = {'classes': tuple(class_names), **metainfo}
+
+        return metainfo
+
+    def full_init(self):
+        """Load annotation file and set ``BaseDataset._fully_initialized`` to
+        True."""
+        super().full_init()
+
+        #  To support the standard OpenMMLab 2.0 annotation format. Generate
+        #  metainfo in internal format from standard metainfo format.
+        if 'categories' in self._metainfo and 'classes' not in self._metainfo:
+            categories = sorted(
+                self._metainfo['categories'], key=lambda x: x['id'])
+            self._metainfo['classes'] = tuple(
+                [cat['category_name'] for cat in categories])
+
+    def __repr__(self):
+        """Print the basic information of the dataset.
+
+        Returns:
+            str: Formatted string.
+        """
+        head = 'Dataset ' + self.__class__.__name__
+        body = []
+        if self._fully_initialized:
+            body.append(f'Number of samples: \t{self.__len__()}')
+        else:
+            body.append("Haven't been initialized")
+
+        if self.CLASSES is not None:
+            body.append(f'Number of categories: \t{len(self.CLASSES)}')
+        else:
+            body.append('The `CLASSES` meta info is not set.')
+
+        body.extend(self.extra_repr())
+
+        if len(self.pipeline.transforms) > 0:
+            body.append('With transforms:')
+            for t in self.pipeline.transforms:
+                body.append(f'    {t}')
+
+        lines = [head] + [' ' * 4 + line for line in body]
+        return '\n'.join(lines)
+
+    def extra_repr(self) -> List[str]:
+        """The extra repr information of the dataset."""
+        body = []
+        body.append(f'Annotation file: \t{self.ann_file}')
+        body.append(f'Prefix of images: \t{self.img_prefix}')
+        return body

mmpl/datasets/builder.py

@@ -0,0 +1,25 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from mmpl.registry import DATASETS
+
+
+def build_dataset(cfg):
+    """Build dataset.
+
+    Examples:
+        >>> from mmpl.datasets import build_dataset
+        >>> mnist_train = build_dataset(
+        ...     dict(type='MNIST', data_prefix='data/mnist/', test_mode=False))
+        >>> print(mnist_train)
+        Dataset MNIST
+            Number of samples:  60000
+            Number of categories:       10
+            Prefix of data:     data/mnist/
+        >>> mnist_test = build_dataset(
+        ...     dict(type='MNIST', data_prefix='data/mnist/', test_mode=True))
+        >>> print(mnist_test)
+        Dataset MNIST
+            Number of samples:  10000
+            Number of categories:       10
+            Prefix of data:     data/mnist/
+    """
+    return DATASETS.build(cfg)

mmpl/datasets/custom.py

@@ -0,0 +1,237 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
+
+from mmengine.fileio import (BaseStorageBackend, get_file_backend,
+                             list_from_file)
+from mmengine.logging import MMLogger
+
+from mmcls.registry import DATASETS
+from .base_dataset import BaseDataset
+
+
+def find_folders(
+    root: str,
+    backend: Optional[BaseStorageBackend] = None
+) -> Tuple[List[str], Dict[str, int]]:
+    """Find classes by folders under a root.
+
+    Args:
+        root (string): root directory of folders
+        backend (BaseStorageBackend | None): The file backend of the root.
+            If None, auto infer backend from the root path. Defaults to None.
+
+    Returns:
+        Tuple[List[str], Dict[str, int]]:
+
+        - folders: The name of sub folders under the root.
+        - folder_to_idx: The map from folder name to class idx.
+    """
+    # Pre-build file backend to prevent verbose file backend inference.
+    backend = backend or get_file_backend(root, enable_singleton=True)
+    folders = list(
+        backend.list_dir_or_file(
+            root,
+            list_dir=True,
+            list_file=False,
+            recursive=False,
+        ))
+    folders.sort()
+    folder_to_idx = {folders[i]: i for i in range(len(folders))}
+    return folders, folder_to_idx
+
+
+def get_samples(
+    root: str,
+    folder_to_idx: Dict[str, int],
+    is_valid_file: Callable,
+    backend: Optional[BaseStorageBackend] = None,
+):
+    """Make dataset by walking all images under a root.
+
+    Args:
+        root (string): root directory of folders
+        folder_to_idx (dict): the map from class name to class idx
+        is_valid_file (Callable): A function that takes path of a file
+            and check if the file is a valid sample file.
+        backend (BaseStorageBackend | None): The file backend of the root.
+            If None, auto infer backend from the root path. Defaults to None.
+
+    Returns:
+        Tuple[list, set]:
+
+        - samples: a list of tuple where each element is (image, class_idx)
+        - empty_folders: The folders don't have any valid files.
+    """
+    samples = []
+    available_classes = set()
+    # Pre-build file backend to prevent verbose file backend inference.
+    backend = backend or get_file_backend(root, enable_singleton=True)
+
+    for folder_name in sorted(list(folder_to_idx.keys())):
+        _dir = backend.join_path(root, folder_name)
+        files = backend.list_dir_or_file(
+            _dir,
+            list_dir=False,
+            list_file=True,
+            recursive=True,
+        )
+        for file in sorted(list(files)):
+            if is_valid_file(file):
+                path = backend.join_path(folder_name, file)
+                item = (path, folder_to_idx[folder_name])
+                samples.append(item)
+                available_classes.add(folder_name)
+
+    empty_folders = set(folder_to_idx.keys()) - available_classes
+
+    return samples, empty_folders
+
+
+@DATASETS.register_module()
+class CustomDataset(BaseDataset):
+    """Custom dataset for classification.
+
+    The dataset supports two kinds of annotation format.
+
+    1. An annotation file is provided, and each line indicates a sample:
+
+       The sample files: ::
+
+           data_prefix/
+           ├── folder_1
+           │   ├── xxx.png
+           │   ├── xxy.png
+           │   └── ...
+           └── folder_2
+               ├── 123.png
+               ├── nsdf3.png
+               └── ...
+
+       The annotation file (the first column is the image path and the second
+       column is the index of category): ::
+
+            folder_1/xxx.png 0
+            folder_1/xxy.png 1
+            folder_2/123.png 5
+            folder_2/nsdf3.png 3
+            ...
+
+       Please specify the name of categories by the argument ``classes``
+       or ``metainfo``.
+
+    2. The samples are arranged in the specific way: ::
+
+           data_prefix/
+           ├── class_x
+           │   ├── xxx.png
+           │   ├── xxy.png
+           │   └── ...
+           │       └── xxz.png
+           └── class_y
+               ├── 123.png
+               ├── nsdf3.png
+               ├── ...
+               └── asd932_.png
+
+    If the ``ann_file`` is specified, the dataset will be generated by the
+    first way, otherwise, try the second way.
+
+    Args:
+        ann_file (str): Annotation file path. Defaults to ''.
+        metainfo (dict, optional): Meta information for dataset, such as class
+            information. Defaults to None.
+        data_root (str): The root directory for ``data_prefix`` and
+            ``ann_file``. Defaults to ''.
+        data_prefix (str | dict): Prefix for the data. Defaults to ''.
+        extensions (Sequence[str]): A sequence of allowed extensions. Defaults
+            to ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif').
+        lazy_init (bool): Whether to load annotation during instantiation.
+            In some cases, such as visualization, only the meta information of
+            the dataset is needed, which is not necessary to load annotation
+            file. ``Basedataset`` can skip load annotations to save time by set
+            ``lazy_init=False``. Defaults to False.
+        **kwargs: Other keyword arguments in :class:`BaseDataset`.
+    """
+
+    def __init__(self,
+                 ann_file: str = '',
+                 metainfo: Optional[dict] = None,
+                 data_root: str = '',
+                 data_prefix: Union[str, dict] = '',
+                 extensions: Sequence[str] = ('.jpg', '.jpeg', '.png', '.ppm',
+                                              '.bmp', '.pgm', '.tif'),
+                 lazy_init: bool = False,
+                 **kwargs):
+        assert (ann_file or data_prefix or data_root), \
+            'One of `ann_file`, `data_root` and `data_prefix` must '\
+            'be specified.'
+
+        self.extensions = tuple(set([i.lower() for i in extensions]))
+
+        super().__init__(
+            # The base class requires string ann_file but this class doesn't
+            ann_file=ann_file,
+            metainfo=metainfo,
+            data_root=data_root,
+            data_prefix=data_prefix,
+            # Force to lazy_init for some modification before loading data.
+            lazy_init=True,
+            **kwargs)
+
+        # Full initialize the dataset.
+        if not lazy_init:
+            self.full_init()
+
+    def _find_samples(self):
+        """find samples from ``data_prefix``."""
+        classes, folder_to_idx = find_folders(self.img_prefix)
+        samples, empty_classes = get_samples(
+            self.img_prefix,
+            folder_to_idx,
+            is_valid_file=self.is_valid_file,
+        )
+
+        if len(samples) == 0:
+            raise RuntimeError(
+                f'Found 0 files in subfolders of: {self.data_prefix}. '
+                f'Supported extensions are: {",".join(self.extensions)}')
+
+        if self.CLASSES is not None:
+            assert len(self.CLASSES) == len(classes), \
+                f"The number of subfolders ({len(classes)}) doesn't match " \
+                f'the number of specified classes ({len(self.CLASSES)}). ' \
+                'Please check the data folder.'
+        else:
+            self._metainfo['classes'] = tuple(classes)
+
+        if empty_classes:
+            logger = MMLogger.get_current_instance()
+            logger.warning(
+                'Found no valid file in the folder '
+                f'{", ".join(empty_classes)}. '
+                f"Supported extensions are: {', '.join(self.extensions)}")
+
+        self.folder_to_idx = folder_to_idx
+
+        return samples
+
+    def load_data_list(self):
+        """Load image paths and gt_labels."""
+        if not self.ann_file:
+            samples = self._find_samples()
+        else:
+            lines = list_from_file(self.ann_file)
+            samples = [x.strip().rsplit(' ', 1) for x in lines]
+
+        # Pre-build file backend to prevent verbose file backend inference.
+        backend = get_file_backend(self.img_prefix, enable_singleton=True)
+        data_list = []
+        for filename, gt_label in samples:
+            img_path = backend.join_path(self.img_prefix, filename)
+            info = {'img_path': img_path, 'gt_label': int(gt_label)}
+            data_list.append(info)
+        return data_list
+
+    def is_valid_file(self, filename: str) -> bool:
+        """Check if a file is a valid sample."""
+        return filename.lower().endswith(self.extensions)

mmpl/datasets/nwpu_ins_dataset.py

@@ -0,0 +1,59 @@
+from typing import List
+
+from mmpl.registry import DATASETS
+from mmdet.datasets.coco import CocoDataset
+
+
+@DATASETS.register_module()
+class NWPUInsSegDataset(CocoDataset):
+    """Dataset for Cityscapes."""
+
+    METAINFO = {
+        'classes': ['airplane', 'ship', 'storage_tank', 'baseball_diamond',
+                    'tennis_court', 'basketball_court', 'ground_track_field',
+                    'harbor', 'bridge', 'vehicle'],
+        'palette': [(220, 20, 60), (255, 0, 0), (0, 0, 142), (0, 0, 70),
+                    (0, 60, 100), (0, 80, 100), (0, 0, 230),
+                    (119, 11, 32), (0, 255, 0), (0, 0, 255)]
+    }
+
+    def filter_data(self) -> List[dict]:
+        """Filter annotations according to filter_cfg.
+
+        Returns:
+            List[dict]: Filtered results.
+        """
+        if self.test_mode:
+            return self.data_list
+
+        if self.filter_cfg is None:
+            return self.data_list
+
+        filter_empty_gt = self.filter_cfg.get('filter_empty_gt', False)
+        min_size = self.filter_cfg.get('min_size', 0)
+
+        # obtain images that contain annotation
+        ids_with_ann = set(data_info['img_id'] for data_info in self.data_list)
+        # obtain images that contain annotations of the required categories
+        ids_in_cat = set()
+        for i, class_id in enumerate(self.cat_ids):
+            ids_in_cat |= set(self.cat_img_map[class_id])
+        # merge the image id sets of the two conditions and use the merged set
+        # to filter out images if self.filter_empty_gt=True
+        ids_in_cat &= ids_with_ann
+
+        valid_data_infos = []
+        for i, data_info in enumerate(self.data_list):
+            img_id = data_info['img_id']
+            width = data_info['width']
+            height = data_info['height']
+            all_is_crowd = all([
+                instance['ignore_flag'] == 1
+                for instance in data_info['instances']
+            ])
+            if filter_empty_gt and (img_id not in ids_in_cat or all_is_crowd):
+                continue
+            if min(width, height) >= min_size:
+                valid_data_infos.append(data_info)
+
+        return valid_data_infos

mmpl/datasets/pl_datamodule.py

@@ -0,0 +1,73 @@
+from mmpl.registry import DATASETS
+import lightning.pytorch as pl
+from torch.utils.data import DataLoader
+from .builder import build_dataset
+from mmengine.registry import FUNCTIONS
+from functools import partial
+
+
+def get_collate_fn(dataloader_cfg):
+    collate_fn_cfg = dataloader_cfg.pop('collate_fn', dict(type='pseudo_collate'))
+    collate_fn_type = collate_fn_cfg.pop('type')
+    collate_fn = FUNCTIONS.get(collate_fn_type)
+    collate_fn = partial(collate_fn, **collate_fn_cfg)  # type: ignore
+    return collate_fn
+
+
+@DATASETS.register_module()
+class PLDataModule(pl.LightningDataModule):
+    def __init__(self,
+                 train_loader=None,
+                 val_loader=None,
+                 test_loader=None,
+                 predict_loader=None,
+                 **kwargs
+                 ):
+        super().__init__()
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        self.test_loader = test_loader
+        self.predict_loader = predict_loader
+        self.train_dataset = None
+        self.val_dataset = None
+        self.test_dataset = None
+        self.predict_dataset = None
+
+    def prepare_data(self):
+        pass
+
+    def setup(self, stage: str):
+        if stage == "fit":
+            dataset_cfg = self.train_loader.pop('dataset')
+            self.train_dataset = build_dataset(dataset_cfg)
+            if self.val_loader is not None:
+                dataset_cfg = self.val_loader.pop('dataset')
+                self.val_dataset = build_dataset(dataset_cfg)
+        if stage == "val":
+            if self.val_loader is not None:
+                dataset_cfg = self.val_loader.pop('dataset')
+                self.val_dataset = build_dataset(dataset_cfg)
+        if stage == "test":
+            if self.test_loader is not None:
+                dataset_cfg = self.test_loader.pop('dataset')
+                self.test_dataset = build_dataset(dataset_cfg)
+        if stage == "predict":
+            if self.predict_loader is not None:
+                dataset_cfg = self.predict_loader.pop('dataset')
+                self.predict_dataset = build_dataset(dataset_cfg)
+
+    def train_dataloader(self):
+        collate_fn = get_collate_fn(self.train_loader)
+        return DataLoader(self.train_dataset, collate_fn=collate_fn, **self.train_loader)
+
+    def val_dataloader(self):
+        collate_fn = get_collate_fn(self.val_loader)
+        return DataLoader(self.val_dataset, collate_fn=collate_fn, **self.val_loader)
+
+    def test_dataloader(self):
+        collate_fn = get_collate_fn(self.test_loader)
+        return DataLoader(self.test_dataset, collate_fn=collate_fn, **self.test_loader)
+
+    def predict_dataloader(self):
+        collate_fn = get_collate_fn(self.predict_loader)
+        return DataLoader(self.predict_dataset, collate_fn=collate_fn, **self.predict_loader)

mmpl/datasets/ssdd_ins_dataset.py

@@ -0,0 +1,54 @@
+from typing import List
+from mmpl.registry import DATASETS
+from mmdet.datasets.coco import CocoDataset
+
+
+@DATASETS.register_module()
+class SSDDInsSegDataset(CocoDataset):
+    """Dataset for Cityscapes."""
+
+    METAINFO = {
+        'classes': ['ship'],
+        'palette': [(0, 0, 255)]
+    }
+
+    def filter_data(self) -> List[dict]:
+        """Filter annotations according to filter_cfg.
+
+        Returns:
+            List[dict]: Filtered results.
+        """
+        # if self.test_mode:
+        #     return self.data_list
+
+        if self.filter_cfg is None:
+            return self.data_list
+
+        filter_empty_gt = self.filter_cfg.get('filter_empty_gt', False)
+        min_size = self.filter_cfg.get('min_size', 0)
+
+        # obtain images that contain annotation
+        ids_with_ann = set(data_info['img_id'] for data_info in self.data_list)
+        # obtain images that contain annotations of the required categories
+        ids_in_cat = set()
+        for i, class_id in enumerate(self.cat_ids):
+            ids_in_cat |= set(self.cat_img_map[class_id])
+        # merge the image id sets of the two conditions and use the merged set
+        # to filter out images if self.filter_empty_gt=True
+        ids_in_cat &= ids_with_ann
+
+        valid_data_infos = []
+        for i, data_info in enumerate(self.data_list):
+            img_id = data_info['img_id']
+            width = data_info['width']
+            height = data_info['height']
+            all_is_crowd = all([
+                instance['ignore_flag'] == 1
+                for instance in data_info['instances']
+            ])
+            if filter_empty_gt and (img_id not in ids_in_cat or all_is_crowd):
+                continue
+            if min(width, height) >= min_size:
+                valid_data_infos.append(data_info)
+
+        return valid_data_infos

mmpl/datasets/utils.py

@@ -0,0 +1,243 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import gzip
+import hashlib
+import os
+import os.path
+import shutil
+import tarfile
+import tempfile
+import urllib.error
+import urllib.request
+import zipfile
+
+from mmengine.fileio import LocalBackend, get_file_backend
+
+__all__ = [
+    'rm_suffix', 'check_integrity', 'download_and_extract_archive',
+    'open_maybe_compressed_file'
+]
+
+
+def rm_suffix(s, suffix=None):
+    if suffix is None:
+        return s[:s.rfind('.')]
+    else:
+        return s[:s.rfind(suffix)]
+
+
+def calculate_md5(fpath: str, chunk_size: int = 1024 * 1024):
+    md5 = hashlib.md5()
+    backend = get_file_backend(fpath, enable_singleton=True)
+    if isinstance(backend, LocalBackend):
+        # Enable chunk update for local file.
+        with open(fpath, 'rb') as f:
+            for chunk in iter(lambda: f.read(chunk_size), b''):
+                md5.update(chunk)
+    else:
+        md5.update(backend.get(fpath))
+    return md5.hexdigest()
+
+
+def check_md5(fpath, md5, **kwargs):
+    return md5 == calculate_md5(fpath, **kwargs)
+
+
+def check_integrity(fpath, md5=None):
+    if not os.path.isfile(fpath):
+        return False
+    if md5 is None:
+        return True
+    return check_md5(fpath, md5)
+
+
+def download_url_to_file(url, dst, hash_prefix=None, progress=True):
+    """Download object at the given URL to a local path.
+
+    Modified from
+    https://pytorch.org/docs/stable/hub.html#torch.hub.download_url_to_file
+
+    Args:
+        url (str): URL of the object to download
+        dst (str): Full path where object will be saved,
+            e.g. ``/tmp/temporary_file``
+        hash_prefix (string, optional): If not None, the SHA256 downloaded
+            file should start with ``hash_prefix``. Defaults to None.
+        progress (bool): whether or not to display a progress bar to stderr.
+            Defaults to True
+    """
+    file_size = None
+    req = urllib.request.Request(url)
+    u = urllib.request.urlopen(req)
+    meta = u.info()
+    if hasattr(meta, 'getheaders'):
+        content_length = meta.getheaders('Content-Length')
+    else:
+        content_length = meta.get_all('Content-Length')
+    if content_length is not None and len(content_length) > 0:
+        file_size = int(content_length[0])
+
+    # We deliberately save it in a temp file and move it after download is
+    # complete. This prevents a local file being overridden by a broken
+    # download.
+    dst = os.path.expanduser(dst)
+    dst_dir = os.path.dirname(dst)
+    f = tempfile.NamedTemporaryFile(delete=False, dir=dst_dir)
+
+    import rich.progress
+    columns = [
+        rich.progress.DownloadColumn(),
+        rich.progress.BarColumn(bar_width=None),
+        rich.progress.TimeRemainingColumn(),
+    ]
+    try:
+        if hash_prefix is not None:
+            sha256 = hashlib.sha256()
+        with rich.progress.Progress(*columns) as pbar:
+            task = pbar.add_task('download', total=file_size, visible=progress)
+            while True:
+                buffer = u.read(8192)
+                if len(buffer) == 0:
+                    break
+                f.write(buffer)
+                if hash_prefix is not None:
+                    sha256.update(buffer)
+                pbar.update(task, advance=len(buffer))
+
+        f.close()
+        if hash_prefix is not None:
+            digest = sha256.hexdigest()
+            if digest[:len(hash_prefix)] != hash_prefix:
+                raise RuntimeError(
+                    'invalid hash value (expected "{}", got "{}")'.format(
+                        hash_prefix, digest))
+        shutil.move(f.name, dst)
+    finally:
+        f.close()
+        if os.path.exists(f.name):
+            os.remove(f.name)
+
+
+def download_url(url, root, filename=None, md5=None):
+    """Download a file from a url and place it in root.
+
+    Args:
+        url (str): URL to download file from.
+        root (str): Directory to place downloaded file in.
+        filename (str | None): Name to save the file under.
+            If filename is None, use the basename of the URL.
+        md5 (str | None): MD5 checksum of the download.
+            If md5 is None, download without md5 check.
+    """
+    root = os.path.expanduser(root)
+    if not filename:
+        filename = os.path.basename(url)
+    fpath = os.path.join(root, filename)
+
+    os.makedirs(root, exist_ok=True)
+
+    if check_integrity(fpath, md5):
+        print(f'Using downloaded and verified file: {fpath}')
+    else:
+        try:
+            print(f'Downloading {url} to {fpath}')
+            download_url_to_file(url, fpath)
+        except (urllib.error.URLError, IOError) as e:
+            if url[:5] == 'https':
+                url = url.replace('https:', 'http:')
+                print('Failed download. Trying https -> http instead.'
+                      f' Downloading {url} to {fpath}')
+                download_url_to_file(url, fpath)
+            else:
+                raise e
+        # check integrity of downloaded file
+        if not check_integrity(fpath, md5):
+            raise RuntimeError('File not found or corrupted.')
+
+
+def _is_tarxz(filename):
+    return filename.endswith('.tar.xz')
+
+
+def _is_tar(filename):
+    return filename.endswith('.tar')
+
+
+def _is_targz(filename):
+    return filename.endswith('.tar.gz')
+
+
+def _is_tgz(filename):
+    return filename.endswith('.tgz')
+
+
+def _is_gzip(filename):
+    return filename.endswith('.gz') and not filename.endswith('.tar.gz')
+
+
+def _is_zip(filename):
+    return filename.endswith('.zip')
+
+
+def extract_archive(from_path, to_path=None, remove_finished=False):
+    if to_path is None:
+        to_path = os.path.dirname(from_path)
+
+    if _is_tar(from_path):
+        with tarfile.open(from_path, 'r') as tar:
+            tar.extractall(path=to_path)
+    elif _is_targz(from_path) or _is_tgz(from_path):
+        with tarfile.open(from_path, 'r:gz') as tar:
+            tar.extractall(path=to_path)
+    elif _is_tarxz(from_path):
+        with tarfile.open(from_path, 'r:xz') as tar:
+            tar.extractall(path=to_path)
+    elif _is_gzip(from_path):
+        to_path = os.path.join(
+            to_path,
+            os.path.splitext(os.path.basename(from_path))[0])
+        with open(to_path, 'wb') as out_f, gzip.GzipFile(from_path) as zip_f:
+            out_f.write(zip_f.read())
+    elif _is_zip(from_path):
+        with zipfile.ZipFile(from_path, 'r') as z:
+            z.extractall(to_path)
+    else:
+        raise ValueError(f'Extraction of {from_path} not supported')
+
+    if remove_finished:
+        os.remove(from_path)
+
+
+def download_and_extract_archive(url,
+                                 download_root,
+                                 extract_root=None,
+                                 filename=None,
+                                 md5=None,
+                                 remove_finished=False):
+    download_root = os.path.expanduser(download_root)
+    if extract_root is None:
+        extract_root = download_root
+    if not filename:
+        filename = os.path.basename(url)
+
+    download_url(url, download_root, filename, md5)
+
+    archive = os.path.join(download_root, filename)
+    print(f'Extracting {archive} to {extract_root}')
+    extract_archive(archive, extract_root, remove_finished)
+
+
+def open_maybe_compressed_file(path: str):
+    """Return a file object that possibly decompresses 'path' on the fly.
+
+    Decompression occurs when argument `path` is a string and ends with '.gz'
+    or '.xz'.
+    """
+    if not isinstance(path, str):
+        return path
+    if path.endswith('.gz'):
+        import gzip
+        return gzip.open(path, 'rb')
+    if path.endswith('.xz'):
+        import lzma
+        return lzma.open(path, 'rb')
+    return open(path, 'rb')

mmpl/datasets/whu_ins_dataset.py

@@ -0,0 +1,54 @@
+from typing import List
+from mmpl.registry import DATASETS
+from mmdet.datasets.coco import CocoDataset
+
+
+@DATASETS.register_module()
+class WHUInsSegDataset(CocoDataset):
+    """Dataset for Cityscapes."""
+
+    METAINFO = {
+        'classes': ['building'],
+        'palette': [(0, 255, 0)]
+    }
+
+    def filter_data(self) -> List[dict]:
+        """Filter annotations according to filter_cfg.
+
+        Returns:
+            List[dict]: Filtered results.
+        """
+        # if self.test_mode:
+        #     return self.data_list
+
+        if self.filter_cfg is None:
+            return self.data_list
+
+        filter_empty_gt = self.filter_cfg.get('filter_empty_gt', False)
+        min_size = self.filter_cfg.get('min_size', 0)
+
+        # obtain images that contain annotation
+        ids_with_ann = set(data_info['img_id'] for data_info in self.data_list)
+        # obtain images that contain annotations of the required categories
+        ids_in_cat = set()
+        for i, class_id in enumerate(self.cat_ids):
+            ids_in_cat |= set(self.cat_img_map[class_id])
+        # merge the image id sets of the two conditions and use the merged set
+        # to filter out images if self.filter_empty_gt=True
+        ids_in_cat &= ids_with_ann
+
+        valid_data_infos = []
+        for i, data_info in enumerate(self.data_list):
+            img_id = data_info['img_id']
+            width = data_info['width']
+            height = data_info['height']
+            all_is_crowd = all([
+                instance['ignore_flag'] == 1
+                for instance in data_info['instances']
+            ])
+            if filter_empty_gt and (img_id not in ids_in_cat or all_is_crowd):
+                continue
+            if min(width, height) >= min_size:
+                valid_data_infos.append(data_info)
+
+        return valid_data_infos

mmpl/engine/__init__.py

@@ -0,0 +1,5 @@
+from .runner import *
+from .logger import *
+from .hooks import *
+from .visualization import *
+from .strategies import *

mmpl/engine/hooks/__init__.py

@@ -0,0 +1,6 @@
+from .builder import PL_HOOKS
+from .pipeline_switch_hook import PipelineSwitchHook
+from .yolov5_param_scheduler_hook import YOLOv5ParamSchedulerHook
+from .ema_hook import EMAHook
+from .param_scheduler_hook import ParamSchedulerHook
+from .visualization_hook import DetVisualizationHook

mmpl/engine/hooks/builder.py

@@ -0,0 +1,31 @@
+import copy
+import inspect
+from typing import List, Union
+
+import torch
+import torch.nn as nn
+import lightning
+
+from mmengine.config import Config, ConfigDict
+from mmengine.device import is_npu_available
+from mmpl.registry import HOOKS
+
+
+def register_pl_hooks() -> List[str]:
+    """Register callbacks in ``lightning.pytorch.callbacks`` to the ``HOOKS`` registry.
+
+    Returns:
+        List[str]: A list of registered callbacks' name.
+    """
+    pl_hooks = []
+    for module_name in dir(lightning.pytorch.callbacks):
+        if module_name.startswith('__'):
+            continue
+        _hook = getattr(lightning.pytorch.callbacks, module_name)
+        if inspect.isclass(_hook) and issubclass(_hook, lightning.pytorch.callbacks.Callback):
+            HOOKS.register_module(module=_hook)
+            pl_hooks.append(module_name)
+    return pl_hooks
+
+
+PL_HOOKS = register_pl_hooks()

mmpl/engine/hooks/ema_hook.py

@@ -0,0 +1,240 @@
+import copy
+import itertools
+import logging
+from typing import Dict, Optional, Any
+
+from lightning import Callback
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from mmengine.logging import print_log
+from mmengine.model import is_model_wrapper
+from mmpl.registry import HOOKS, MODELS
+
+
+
+@HOOKS.register_module()
+class EMAHook(Callback):
+    """A Hook to apply Exponential Moving Average (EMA) on the model during
+    training.
+
+    Note:
+        - EMAHook takes priority over CheckpointHook.
+        - The original model parameters are actually saved in ema field after
+          train.
+        - ``begin_iter`` and ``begin_epoch`` cannot be set at the same time.
+
+    Args:
+        ema_type (str): The type of EMA strategy to use. You can find the
+            supported strategies in :mod:`mmengine.model.averaged_model`.
+            Defaults to 'ExponentialMovingAverage'.
+        strict_load (bool): Whether to strictly enforce that the keys of
+            ``state_dict`` in checkpoint match the keys returned by
+            ``self.module.state_dict``. Defaults to False.
+            Changed in v0.3.0.
+        begin_iter (int): The number of iteration to enable ``EMAHook``.
+            Defaults to 0.
+        begin_epoch (int): The number of epoch to enable ``EMAHook``.
+            Defaults to 0.
+        **kwargs: Keyword arguments passed to subclasses of
+            :obj:`BaseAveragedModel`
+    """
+
+    priority = 'NORMAL'
+
+    def __init__(self,
+                 ema_type: str = 'ExponentialMovingAverage',
+                 strict_load: bool = False,
+                 begin_iter: int = 0,
+                 begin_epoch: int = 0,
+                 **kwargs):
+        self.strict_load = strict_load
+        self.ema_cfg = dict(type=ema_type, **kwargs)
+        assert not (begin_iter != 0 and begin_epoch != 0), (
+            '`begin_iter` and `begin_epoch` should not be both set.')
+        assert begin_iter >= 0, (
+            '`begin_iter` must larger than or equal to 0, '
+            f'but got begin_iter: {begin_iter}')
+        assert begin_epoch >= 0, (
+            '`begin_epoch` must larger than or equal to 0, '
+            f'but got begin_epoch: {begin_epoch}')
+        self.begin_iter = begin_iter
+        self.begin_epoch = begin_epoch
+        # If `begin_epoch` and `begin_iter` are not set, `EMAHook` will be
+        # enabled at 0 iteration.
+        self.enabled_by_epoch = self.begin_epoch > 0
+
+    def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """Create an ema copy of the model.
+
+        Args:
+            runner (Runner): The runner of the training process.
+        """
+        model = pl_module
+        if is_model_wrapper(model):
+            model = model.module
+        self.src_model = model
+        self.ema_model = MODELS.build(
+            self.ema_cfg, default_args=dict(model=self.src_model))
+
+    def on_train_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """Check the begin_epoch/iter is smaller than max_epochs/iters.
+
+        Args:
+            runner (Runner): The runner of the training process.
+        """
+        if self.enabled_by_epoch:
+            assert self.begin_epoch <= trainer.max_epochs, (
+                'self.begin_epoch should be smaller than or equal to '
+                f'runner.max_epochs: {trainer.max_epochs}, but got '
+                f'begin_epoch: {self.begin_epoch}')
+        else:
+            assert self.begin_iter <= trainer.max_steps or self.begin_iter <= trainer.max_epochs * len(trainer.train_dataloader), (
+                'self.begin_iter should be smaller than or equal to '
+                f'runner.max_iters: {trainer.max_steps}, but got '
+                f'begin_iter: {self.begin_iter}')
+
+    def on_train_batch_end(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", outputs: STEP_OUTPUT, batch: Any, batch_idx: int
+    ) -> None:
+        """Update ema parameter.
+
+        Args:
+            runner (Runner): The runner of the training process.
+            batch_idx (int): The index of the current batch in the train loop.
+            data_batch (Sequence[dict], optional): Data from dataloader.
+                Defaults to None.
+            outputs (dict, optional): Outputs from model. Defaults to None.
+        """
+        if self._ema_started(trainer):
+            self.ema_model.update_parameters(self.src_model)
+        else:
+            ema_params = self.ema_model.module.state_dict()
+            src_params = self.src_model.state_dict()
+            for k, p in ema_params.items():
+                p.data.copy_(src_params[k].data)
+
+    def on_validation_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """We load parameter values from ema model to source model before
+        validation.
+
+        Args:
+            runner (Runner): The runner of the training process.
+        """
+        self._swap_ema_parameters()
+
+    def on_validation_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """We recover source model's parameter from ema model after validation.
+
+        Args:
+            runner (Runner): The runner of the validation process.
+            metrics (Dict[str, float], optional): Evaluation results of all
+                metrics on validation dataset. The keys are the names of the
+                metrics, and the values are corresponding results.
+        """
+        self._swap_ema_parameters()
+
+    def on_test_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """We load parameter values from ema model to source model before test.
+
+        Args:
+            runner (Runner): The runner of the training process.
+        """
+        self._swap_ema_parameters()
+
+    def on_test_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """We recover source model's parameter from ema model after test.
+
+        Args:
+            runner (Runner): The runner of the testing process.
+            metrics (Dict[str, float], optional): Evaluation results of all
+                metrics on test dataset. The keys are the names of the
+                metrics, and the values are corresponding results.
+        """
+        self._swap_ema_parameters()
+
+    def on_save_checkpoint(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", checkpoint: Dict[str, Any]
+    ) -> None:
+        """Save ema parameters to checkpoint.
+
+        Args:
+            runner (Runner): The runner of the testing process.
+        """
+        checkpoint['ema_state_dict'] = self.ema_model.state_dict()
+        # Save ema parameters to the source model's state dict so that we
+        # can directly load the averaged model weights for deployment.
+        # Swapping the state_dict key-values instead of swapping model
+        # parameters because the state_dict is a shallow copy of model
+        # parameters.
+        self._swap_ema_state_dict(checkpoint)
+
+    def on_load_checkpoint(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", checkpoint: Dict[str, Any]
+    ) -> None:
+        """Resume ema parameters from checkpoint.
+
+        Args:
+            runner (Runner): The runner of the testing process.
+        """
+        from mmengine.runner.checkpoint import load_state_dict
+        if 'ema_state_dict' in checkpoint and not trainer._checkpoint_connector._loaded_checkpoint:
+            # The original model parameters are actually saved in ema
+            # field swap the weights back to resume ema state.
+            self._swap_ema_state_dict(checkpoint)
+            self.ema_model.load_state_dict(
+                checkpoint['ema_state_dict'], strict=self.strict_load)
+
+        # Support load checkpoint without ema state dict.
+        else:
+            if not trainer._checkpoint_connector._loaded_checkpoint:
+                print_log(
+                    'There is no `ema_state_dict` in checkpoint. '
+                    '`EMAHook` will make a copy of `state_dict` as the '
+                    'initial `ema_state_dict`', 'current', logging.WARNING)
+            load_state_dict(
+                self.ema_model.module,
+                copy.deepcopy(checkpoint['state_dict']),
+                strict=self.strict_load)
+
+    def _swap_ema_parameters(self) -> None:
+        """Swap the parameter of model with ema_model."""
+        avg_param = (
+            itertools.chain(self.ema_model.module.parameters(),
+                            self.ema_model.module.buffers())
+            if self.ema_model.update_buffers else
+            self.ema_model.module.parameters())
+        src_param = (
+            itertools.chain(self.src_model.parameters(),
+                            self.src_model.buffers())
+            if self.ema_model.update_buffers else self.src_model.parameters())
+        for p_avg, p_src in zip(avg_param, src_param):
+            tmp = p_avg.data.clone()
+            p_avg.data.copy_(p_src.data)
+            p_src.data.copy_(tmp)
+
+    def _swap_ema_state_dict(self, checkpoint):
+        """Swap the state dict values of model with ema_model."""
+        model_state = checkpoint['state_dict']
+        ema_state = checkpoint['ema_state_dict']
+        for k in ema_state:
+            if k[:7] == 'module.':
+                tmp = ema_state[k]
+                ema_state[k] = model_state[k[7:]]
+                model_state[k[7:]] = tmp
+
+    def _ema_started(self, trainer) -> bool:
+        """Whether ``EMAHook`` has been initialized at current iteration or
+        epoch.
+
+        :attr:`ema_model` will be initialized when ``runner.iter`` or
+        ``runner.epoch`` is greater than ``self.begin`` for the first time.
+
+        Args:
+            runner (Runner): Runner of the training, validation process.
+
+        Returns:
+            bool: Whether ``EMAHook`` has been initialized.
+        """
+        if self.enabled_by_epoch:
+            return trainer.current_epoch + 1 >= self.begin_epoch
+        else:
+            return trainer.global_step + 1 >= self.begin_iter

mmpl/engine/hooks/param_scheduler_hook.py

@@ -0,0 +1,128 @@
+from typing import Dict, Optional, Union, Any
+
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from mmengine.optim import _ParamScheduler
+from mmpl.registry import HOOKS
+from mmengine.utils import is_list_of
+from lightning import Callback
+
+DATA_BATCH = Optional[Union[dict, tuple, list]]
+
+
+@HOOKS.register_module()
+class ParamSchedulerHook(Callback):
+    """A hook to update some hyper-parameters in optimizer, e.g., learning rate
+    and momentum."""
+
+    priority = 'LOW'
+
+    def on_train_batch_end(
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", outputs: STEP_OUTPUT, batch: Any, batch_idx: int
+    ) -> None:
+        """Call step function for each scheduler after each training iteration.
+
+        Args:
+            runner (Runner): The runner of the training process.
+            batch_idx (int): The index of the current batch in the train loop.
+            data_batch (dict or tuple or list, optional): Data from dataloader.
+                In order to keep this interface consistent with other hooks,
+                we keep ``data_batch`` here.
+            outputs (dict, optional): Outputs from model.
+                In order to keep this interface consistent with other hooks, we
+                keep ``data_batch`` here.
+        """
+        param_schedulers = pl_module.lr_schedulers()
+        if param_schedulers is None:
+            return
+
+        def step(param_schedulers):
+            assert isinstance(param_schedulers, list)
+            for scheduler in param_schedulers:
+                if not scheduler.by_epoch:
+                    scheduler.step()
+        if isinstance(param_schedulers, _ParamScheduler):
+            param_schedulers = [param_schedulers]
+        if isinstance(param_schedulers, list):
+            step(param_schedulers)
+        elif isinstance(param_schedulers, dict):
+            for param_schedulers in param_schedulers.values():
+                step(param_schedulers)
+        else:
+            raise TypeError(
+                'runner.param_schedulers should be list of ParamScheduler or '
+                'a dict containing list of ParamScheduler, '
+                f'but got {param_schedulers}')
+
+    def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """Call step function for each scheduler after each training epoch.
+
+        Args:
+            runner (Runner): The runner of the training process.
+        """
+        param_schedulers = pl_module.lr_schedulers()
+        if param_schedulers is None:
+            return
+
+        def step(param_schedulers):
+            assert isinstance(param_schedulers, list)
+            for scheduler in param_schedulers:
+                if scheduler.by_epoch:
+                    scheduler.step()
+        if isinstance(param_schedulers, _ParamScheduler):
+            param_schedulers = [param_schedulers]
+        if isinstance(param_schedulers, list):
+            step(param_schedulers)
+        elif isinstance(param_schedulers, dict):
+            for param_schedulers in param_schedulers.values():
+                step(param_schedulers)
+        else:
+            raise TypeError(
+                'runner.param_schedulers should be list of ParamScheduler or '
+                'a dict containing list of ParamScheduler, '
+                f'but got {param_schedulers}')
+
+    def on_validation_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """Call step function for each scheduler which has attribute
+        ``need_val_args`` after each validation epoch.
+
+        Args:
+            runner (Runner): The runner of the validation process.
+            metrics (Dict[str, float], optional): Evaluation results of all
+                metrics on validation dataset. The keys are the names of the
+                metrics, and the values are corresponding results.
+
+        Note:
+            if ``runner.param_schedulers`` is not built before,
+            the hook ``after_val_epoch`` will be skipped.
+        """
+        param_schedulers = pl_module.lr_schedulers()
+        if param_schedulers is None:
+            return
+
+        # avoid counting scheduler._global_step
+        # it has counted in after_train_* hook
+        metrics = trainer.callback_metrics
+        if metrics is None:
+            return
+
+        def step(param_schedulers):
+            # check param_schedulers is list and built
+            if not is_list_of(param_schedulers, _ParamScheduler):
+                return
+
+            for scheduler in param_schedulers:
+                if (scheduler.by_epoch
+                        and getattr(scheduler, 'need_val_args', False)):
+                    scheduler.step(metrics)
+        if isinstance(param_schedulers, _ParamScheduler):
+            param_schedulers = [param_schedulers]
+        if isinstance(param_schedulers, list):
+            step(param_schedulers)
+        elif isinstance(param_schedulers, dict):
+            for param_schedulers in param_schedulers.values():
+                step(param_schedulers)
+        else:
+            raise TypeError(
+                'runner.param_schedulers should be list of ParamScheduler or '
+                'a dict containing list of ParamScheduler, '
+                f'but got {param_schedulers}')

mmpl/engine/hooks/pipeline_switch_hook.py

@@ -0,0 +1,41 @@
+from mmcv.transforms import Compose
+from mmpl.registry import HOOKS
+from lightning.pytorch.callbacks import Callback
+
+
+@HOOKS.register_module()
+class PipelineSwitchHook(Callback):
+    """Switch data pipeline at switch_epoch.
+
+    Args:
+        switch_epoch (int): switch pipeline at this epoch.
+        switch_pipeline (list[dict]): the pipeline to switch to.
+    """
+
+    def __init__(self, switch_epoch, switch_pipeline):
+        self.switch_epoch = switch_epoch
+        self.switch_pipeline = switch_pipeline
+        self._restart_dataloader = False
+
+    def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """switch pipeline."""
+        epoch = trainer.current_epoch
+        train_loader = trainer.train_dataloader
+        if epoch == self.switch_epoch:
+            if trainer.local_rank == 0:
+                print('Switch pipeline now!')
+            # The dataset pipeline cannot be updated when persistent_workers
+            # is True, so we need to force the dataloader's multi-process
+            # restart. This is a very hacky approach.
+            train_loader.dataset.pipeline = Compose(self.switch_pipeline)
+            if hasattr(train_loader, 'persistent_workers'
+                       ) and train_loader.persistent_workers is True:
+                train_loader._DataLoader__initialized = False
+                train_loader._iterator = None
+                self._restart_dataloader = True
+
+        else:
+            # Once the restart is complete, we need to restore
+            # the initialization flag.
+            if self._restart_dataloader:
+                train_loader._DataLoader__initialized = True

mmpl/engine/hooks/ppyoloe_param_scheduler_hook.py

@@ -0,0 +1,96 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import math
+from typing import Optional
+
+from mmengine.hooks import ParamSchedulerHook
+from mmengine.runner import Runner
+
+from mmyolo.registry import HOOKS
+
+
+@HOOKS.register_module()
+class PPYOLOEParamSchedulerHook(ParamSchedulerHook):
+    """A hook to update learning rate and momentum in optimizer of PPYOLOE. We
+    use this hook to implement adaptive computation for `warmup_total_iters`,
+    which is not possible with the built-in ParamScheduler in mmyolo.
+
+    Args:
+        warmup_min_iter (int): Minimum warmup iters. Defaults to 1000.
+        start_factor (float): The number we multiply learning rate in the
+            first epoch. The multiplication factor changes towards end_factor
+            in the following epochs. Defaults to 0.
+        warmup_epochs (int): Epochs for warmup. Defaults to 5.
+        min_lr_ratio (float): Minimum learning rate ratio.
+        total_epochs (int): In PPYOLOE, `total_epochs` is set to
+            training_epochs x 1.2. Defaults to 360.
+    """
+    priority = 9
+
+    def __init__(self,
+                 warmup_min_iter: int = 1000,
+                 start_factor: float = 0.,
+                 warmup_epochs: int = 5,
+                 min_lr_ratio: float = 0.0,
+                 total_epochs: int = 360):
+
+        self.warmup_min_iter = warmup_min_iter
+        self.start_factor = start_factor
+        self.warmup_epochs = warmup_epochs
+        self.min_lr_ratio = min_lr_ratio
+        self.total_epochs = total_epochs
+
+        self._warmup_end = False
+        self._base_lr = None
+
+    def before_train(self, runner: Runner):
+        """Operations before train.
+
+        Args:
+            runner (Runner): The runner of the training process.
+        """
+        optimizer = runner.optim_wrapper.optimizer
+        for group in optimizer.param_groups:
+            # If the param is never be scheduled, record the current value
+            # as the initial value.
+            group.setdefault('initial_lr', group['lr'])
+
+        self._base_lr = [
+            group['initial_lr'] for group in optimizer.param_groups
+        ]
+        self._min_lr = [i * self.min_lr_ratio for i in self._base_lr]
+
+    def before_train_iter(self,
+                          runner: Runner,
+                          batch_idx: int,
+                          data_batch: Optional[dict] = None):
+        """Operations before each training iteration.
+
+        Args:
+            runner (Runner): The runner of the training process.
+            batch_idx (int): The index of the current batch in the train loop.
+            data_batch (dict or tuple or list, optional): Data from dataloader.
+        """
+        cur_iters = runner.iter
+        optimizer = runner.optim_wrapper.optimizer
+        dataloader_len = len(runner.train_dataloader)
+
+        # The minimum warmup is self.warmup_min_iter
+        warmup_total_iters = max(
+            round(self.warmup_epochs * dataloader_len), self.warmup_min_iter)
+
+        if cur_iters <= warmup_total_iters:
+            # warm up
+            alpha = cur_iters / warmup_total_iters
+            factor = self.start_factor * (1 - alpha) + alpha
+
+            for group_idx, param in enumerate(optimizer.param_groups):
+                param['lr'] = self._base_lr[group_idx] * factor
+        else:
+            for group_idx, param in enumerate(optimizer.param_groups):
+                total_iters = self.total_epochs * dataloader_len
+                lr = self._min_lr[group_idx] + (
+                    self._base_lr[group_idx] -
+                    self._min_lr[group_idx]) * 0.5 * (
+                        math.cos((cur_iters - warmup_total_iters) * math.pi /
+                                 (total_iters - warmup_total_iters)) + 1.0)
+                param['lr'] = lr

mmpl/engine/hooks/switch_to_deploy_hook.py

@@ -0,0 +1,21 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+from mmengine.hooks import Hook
+from mmengine.runner import Runner
+
+from mmyolo.registry import HOOKS
+from mmyolo.utils import switch_to_deploy
+
+
+@HOOKS.register_module()
+class SwitchToDeployHook(Hook):
+    """Switch to deploy mode before testing.
+
+    This hook converts the multi-channel structure of the training network
+    (high performance) to the one-way structure of the testing network (fast
+    speed and  memory saving).
+    """
+
+    def before_test_epoch(self, runner: Runner):
+        """Switch to deploy mode before testing."""
+        switch_to_deploy(runner.model)

mmpl/engine/hooks/visualization_hook.py

@@ -0,0 +1,199 @@
+import os.path as osp
+import warnings
+from typing import Optional, Sequence, Any
+
+import mmcv
+from lightning import Callback
+from mmengine.fileio import get
+from mmengine.hooks import Hook
+from mmengine.runner import Runner
+from mmengine.utils import mkdir_or_exist
+from mmengine.visualization import Visualizer
+
+from mmpl.registry import HOOKS
+from mmdet.structures import DetDataSample
+
+
+@HOOKS.register_module()
+class DetVisualizationHook(Callback):
+    """Detection Visualization Hook. Used to visualize validation and testing
+    process prediction results.
+
+    In the testing phase:
+
+    1. If ``show`` is True, it means that only the prediction results are
+        visualized without storing data, so ``vis_backends`` needs to
+        be excluded.
+    2. If ``test_out_dir`` is specified, it means that the prediction results
+        need to be saved to ``test_out_dir``. In order to avoid vis_backends
+        also storing data, so ``vis_backends`` needs to be excluded.
+    3. ``vis_backends`` takes effect if the user does not specify ``show``
+        and `test_out_dir``. You can set ``vis_backends`` to WandbVisBackend or
+        TensorboardVisBackend to store the prediction result in Wandb or
+        Tensorboard.
+
+    Args:
+        draw (bool): whether to draw prediction results. If it is False,
+            it means that no drawing will be done. Defaults to False.
+        interval (int): The interval of visualization. Defaults to 50.
+        score_thr (float): The threshold to visualize the bboxes
+            and masks. Defaults to 0.3.
+        show (bool): Whether to display the drawn image. Default to False.
+        wait_time (float): The interval of show (s). Defaults to 0.
+        test_out_dir (str, optional): directory where painted images
+            will be saved in testing process.
+        backend_args (dict, optional): Arguments to instantiate the
+            corresponding backend. Defaults to None.
+    """
+
+    def __init__(self,
+                 draw: bool = False,
+                 interval: int = 50,
+                 score_thr: float = 0.3,
+                 show: bool = False,
+                 wait_time: float = 0.,
+                 test_out_dir: Optional[str] = None,
+                 backend_args: dict = None):
+        self._visualizer: Visualizer = Visualizer.get_current_instance()
+        self.interval = interval
+        self.score_thr = score_thr
+        self.show = show
+        if self.show:
+            # No need to think about vis backends.
+            self._visualizer._vis_backends = {}
+            warnings.warn('The show is True, it means that only '
+                          'the prediction results are visualized '
+                          'without storing data, so vis_backends '
+                          'needs to be excluded.')
+
+        self.wait_time = wait_time
+        self.backend_args = backend_args
+        self.draw = draw
+        self.test_out_dir = test_out_dir
+        self._test_index = 0
+
+    def after_val_iter(self, runner: Runner, batch_idx: int, data_batch: dict,
+                       outputs: Sequence[DetDataSample]) -> None:
+        """Run after every ``self.interval`` validation iterations.
+
+        Args:
+            runner (:obj:`Runner`): The runner of the validation process.
+            batch_idx (int): The index of the current batch in the val loop.
+            data_batch (dict): Data from dataloader.
+            outputs (Sequence[:obj:`DetDataSample`]]): A batch of data samples
+                that contain annotations and predictions.
+        """
+        if self.draw is False:
+            return
+
+        # There is no guarantee that the same batch of images
+        # is visualized for each evaluation.
+        total_curr_iter = runner.iter + batch_idx
+
+        # Visualize only the first data
+        img_path = outputs[0].img_path
+        img_bytes = get(img_path, backend_args=self.backend_args)
+        img = mmcv.imfrombytes(img_bytes, channel_order='rgb')
+
+        if total_curr_iter % self.interval == 0:
+            self._visualizer.add_datasample(
+                osp.basename(img_path) if self.show else 'val_img',
+                img,
+                data_sample=outputs[0],
+                show=self.show,
+                wait_time=self.wait_time,
+                pred_score_thr=self.score_thr,
+                step=total_curr_iter)
+
+    def after_test_iter(self, runner: Runner, batch_idx: int, data_batch: dict,
+                        outputs: Sequence[DetDataSample]) -> None:
+        """Run after every testing iterations.
+
+        Args:
+            runner (:obj:`Runner`): The runner of the testing process.
+            batch_idx (int): The index of the current batch in the val loop.
+            data_batch (dict): Data from dataloader.
+            outputs (Sequence[:obj:`DetDataSample`]): A batch of data samples
+                that contain annotations and predictions.
+        """
+        if self.draw is False:
+            return
+
+        if self.test_out_dir is not None:
+            self.test_out_dir = osp.join(runner.work_dir, runner.timestamp,
+                                         self.test_out_dir)
+            mkdir_or_exist(self.test_out_dir)
+
+        for data_sample in outputs:
+            self._test_index += 1
+
+            img_path = data_sample.img_path
+            img_bytes = get(img_path, backend_args=self.backend_args)
+            img = mmcv.imfrombytes(img_bytes, channel_order='rgb')
+
+            out_file = None
+            if self.test_out_dir is not None:
+                out_file = osp.basename(img_path)
+                out_file = osp.join(self.test_out_dir, out_file)
+
+            self._visualizer.add_datasample(
+                osp.basename(img_path) if self.show else 'test_img',
+                img,
+                data_sample=data_sample,
+                show=self.show,
+                wait_time=self.wait_time,
+                pred_score_thr=self.score_thr,
+                out_file=out_file,
+                step=self._test_index)
+
+    def on_predict_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        # if hasattr(trainer.datamodule, f'predict_dataset'):
+        #     dataset = getattr(trainer.datamodule, f'predict_dataset')
+        #     if hasattr(dataset, 'metainfo') and hasattr(self._visualizer, 'dataset_meta'):
+        #         self._visualizer.dataset_meta = dataset.metainfo
+        if self.test_out_dir is not None:
+            self.test_out_dir = osp.join(trainer.default_root_dir, self.test_out_dir)
+            mkdir_or_exist(self.test_out_dir)
+
+    def on_predict_batch_end(
+        self,
+        trainer: "pl.Trainer",
+        pl_module: "pl.LightningModule",
+        outputs: Any,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        """Run after every testing iterations.
+
+        Args:
+            runner (:obj:`Runner`): The runner of the testing process.
+            batch_idx (int): The index of the current batch in the val loop.
+            data_batch (dict): Data from dataloader.
+            outputs (Sequence[:obj:`DetDataSample`]): A batch of data samples
+                that contain annotations and predictions.
+        """
+        if self.draw is False:
+            return
+
+        for data_sample in outputs:
+            self._test_index += 1
+
+            img_path = data_sample.img_path
+            img_bytes = get(img_path, backend_args=self.backend_args)
+            img = mmcv.imfrombytes(img_bytes, channel_order='rgb')
+
+            out_file = None
+            if self.test_out_dir is not None:
+                out_file = osp.basename(img_path)
+                out_file = osp.join(self.test_out_dir, out_file)
+
+            self._visualizer.add_datasample(
+                osp.basename(img_path) if self.show else 'test_img',
+                img,
+                data_sample=data_sample,
+                show=self.show,
+                wait_time=self.wait_time,
+                pred_score_thr=self.score_thr,
+                out_file=out_file,
+                step=self._test_index)

mmpl/engine/hooks/yolov5_param_scheduler_hook.py

@@ -0,0 +1,111 @@
+import math
+from typing import Optional
+import numpy as np
+from typing import Dict, Optional, Union
+from mmengine.registry import HOOKS
+from .param_scheduler_hook import ParamSchedulerHook
+
+DATA_BATCH = Optional[Union[dict, tuple, list]]
+
+
+def linear_fn(lr_factor: float, max_epochs: int):
+    """Generate linear function."""
+    return lambda x: (1 - x / max_epochs) * (1.0 - lr_factor) + lr_factor
+
+
+def cosine_fn(lr_factor: float, max_epochs: int):
+    """Generate cosine function."""
+    return lambda x: (
+        (1 - math.cos(x * math.pi / max_epochs)) / 2) * (lr_factor - 1) + 1
+
+
+@HOOKS.register_module()
+class YOLOv5ParamSchedulerHook(ParamSchedulerHook):
+    """A hook to update learning rate and momentum in optimizer of YOLOv5."""
+    priority = 9
+
+    scheduler_maps = {'linear': linear_fn, 'cosine': cosine_fn}
+
+    def __init__(self,
+                 scheduler_type: str = 'linear',
+                 lr_factor: float = 0.01,
+                 max_epochs: int = 300,
+                 warmup_epochs: int = 3,
+                 warmup_bias_lr: float = 0.1,
+                 warmup_momentum: float = 0.8,
+                 warmup_mim_iter: int = 500,
+                 **kwargs):
+
+        assert scheduler_type in self.scheduler_maps
+
+        self.warmup_epochs = warmup_epochs
+        self.warmup_bias_lr = warmup_bias_lr
+        self.warmup_momentum = warmup_momentum
+        self.warmup_mim_iter = warmup_mim_iter
+
+        kwargs.update({'lr_factor': lr_factor, 'max_epochs': max_epochs})
+        self.scheduler_fn = self.scheduler_maps[scheduler_type](**kwargs)
+
+        self._warmup_end = False
+        self._base_lr = None
+        self._base_momentum = None
+
+
+    def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        optimizer = trainer.optimizers[0]
+        for group in optimizer.param_groups:
+            # If the param is never be scheduled, record the current value
+            # as the initial value.
+            group.setdefault('initial_lr', group['lr'])
+            group.setdefault('initial_momentum', group.get('momentum', -1))
+
+        self._base_lr = [
+            group['initial_lr'] for group in optimizer.param_groups
+        ]
+        self._base_momentum = [
+            group['initial_momentum'] for group in optimizer.param_groups
+        ]
+
+    def on_before_backward(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", loss) -> None:
+        cur_iters = trainer.global_step
+        cur_epoch = trainer.current_epoch
+        optimizer = trainer.optimizers[0]
+
+        # The minimum warmup is self.warmup_mim_iter
+        warmup_total_iters = max(
+            round(self.warmup_epochs * len(trainer.train_dataloader)),
+            self.warmup_mim_iter)
+
+        if cur_iters <= warmup_total_iters:
+            xp = [0, warmup_total_iters]
+            for group_idx, param in enumerate(optimizer.param_groups):
+                if group_idx == 2:
+                    # bias learning rate will be handled specially
+                    yp = [
+                        self.warmup_bias_lr,
+                        self._base_lr[group_idx] * self.scheduler_fn(cur_epoch)
+                    ]
+                else:
+                    yp = [
+                        0.0,
+                        self._base_lr[group_idx] * self.scheduler_fn(cur_epoch)
+                    ]
+                param['lr'] = np.interp(cur_iters, xp, yp)
+
+                if 'momentum' in param:
+                    param['momentum'] = np.interp(
+                        cur_iters, xp,
+                        [self.warmup_momentum, self._base_momentum[group_idx]])
+        else:
+            self._warmup_end = True
+
+    def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        if not self._warmup_end:
+            return
+
+        cur_epoch = trainer.current_epoch
+        optimizer = trainer.optimizers[0]
+        for group_idx, param in enumerate(optimizer.param_groups):
+            param['lr'] = self._base_lr[group_idx] * self.scheduler_fn(
+                cur_epoch)
+

mmpl/engine/hooks/yolox_mode_switch_hook.py

@@ -0,0 +1,54 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import copy
+from typing import Sequence
+
+from mmengine.hooks import Hook
+from mmengine.model import is_model_wrapper
+from mmengine.runner import Runner
+
+from mmyolo.registry import HOOKS
+
+
+@HOOKS.register_module()
+class YOLOXModeSwitchHook(Hook):
+    """Switch the mode of YOLOX during training.
+
+    This hook turns off the mosaic and mixup data augmentation and switches
+    to use L1 loss in bbox_head.
+
+    Args:
+        num_last_epochs (int): The number of latter epochs in the end of the
+            training to close the data augmentation and switch to L1 loss.
+            Defaults to 15.
+    """
+
+    def __init__(self,
+                 num_last_epochs: int = 15,
+                 new_train_pipeline: Sequence[dict] = None):
+        self.num_last_epochs = num_last_epochs
+        self.new_train_pipeline_cfg = new_train_pipeline
+
+    def before_train_epoch(self, runner: Runner):
+        """Close mosaic and mixup augmentation and switches to use L1 loss."""
+        epoch = runner.epoch
+        model = runner.model
+        if is_model_wrapper(model):
+            model = model.module
+
+        if (epoch + 1) == runner.max_epochs - self.num_last_epochs:
+            runner.logger.info(f'New Pipeline: {self.new_train_pipeline_cfg}')
+
+            train_dataloader_cfg = copy.deepcopy(runner.cfg.train_dataloader)
+            train_dataloader_cfg.dataset.pipeline = self.new_train_pipeline_cfg
+            # Note: Why rebuild the dataset?
+            # When build_dataloader will make a deep copy of the dataset,
+            # it will lead to potential risks, such as the global instance
+            # object FileClient data is disordered.
+            # This problem needs to be solved in the future.
+            new_train_dataloader = Runner.build_dataloader(
+                train_dataloader_cfg)
+            runner.train_loop.dataloader = new_train_dataloader
+
+            runner.logger.info('recreate the dataloader!')
+            runner.logger.info('Add additional bbox reg loss now!')
+            model.bbox_head.use_bbox_aux = True

mmpl/engine/logger/__init__.py

@@ -0,0 +1 @@
+from .builder import PL_LOGGERS

mmpl/engine/logger/builder.py

@@ -0,0 +1,112 @@
+import copy
+import inspect
+from typing import List, Union
+
+import torch
+import torch.nn as nn
+import lightning
+
+from mmengine.config import Config, ConfigDict
+from mmengine.device import is_npu_available
+from mmpl.registry import LOGGERS
+
+
+def register_pl_loggers() -> List[str]:
+    """Register loggers in ``lightning.pytorch.loggers`` to the ``LOGGERS`` registry.
+
+    Returns:
+        List[str]: A list of registered optimizers' name.
+    """
+    pl_loggers = []
+    for module_name in dir(lightning.pytorch.loggers):
+        if module_name.startswith('__'):
+            continue
+        _logger = getattr(lightning.pytorch.loggers, module_name)
+        if inspect.isclass(_logger) and issubclass(_logger, lightning.pytorch.loggers.logger.Logger):
+            LOGGERS.register_module(module=_logger)
+            pl_loggers.append(module_name)
+    return pl_loggers
+
+
+PL_LOGGERS = register_pl_loggers()
+
+
+def register_dadaptation_optimizers() -> List[str]:
+    """Register optimizers in ``dadaptation`` to the ``OPTIMIZERS`` registry.
+
+    Returns:
+        List[str]: A list of registered optimizers' name.
+    """
+    dadaptation_optimizers = []
+    try:
+        import dadaptation
+    except ImportError:
+        pass
+    else:
+        for module_name in ['DAdaptAdaGrad', 'DAdaptAdam', 'DAdaptSGD']:
+            _optim = getattr(dadaptation, module_name)
+            if inspect.isclass(_optim) and issubclass(_optim,
+                                                      torch.optim.Optimizer):
+                OPTIMIZERS.register_module(module=_optim)
+                dadaptation_optimizers.append(module_name)
+    return dadaptation_optimizers
+
+
+# DADAPTATION_OPTIMIZERS = register_dadaptation_optimizers()
+
+
+def register_lion_optimizers() -> List[str]:
+    """Register Lion optimizer to the ``OPTIMIZERS`` registry.
+
+    Returns:
+        List[str]: A list of registered optimizers' name.
+    """
+    optimizers = []
+    try:
+        from lion_pytorch import Lion
+    except ImportError:
+        pass
+    else:
+        OPTIMIZERS.register_module(module=Lion)
+        optimizers.append('Lion')
+    return optimizers
+
+
+# LION_OPTIMIZERS = register_lion_optimizers()
+
+
+def build_optim_wrapper(model: nn.Module,
+                        cfg: Union[dict, Config, ConfigDict]):
+    """Build function of OptimWrapper.
+
+    If ``constructor`` is set in the ``cfg``, this method will build an
+    optimizer wrapper constructor, and use optimizer wrapper constructor to
+    build the optimizer wrapper. If ``constructor`` is not set, the
+    ``DefaultOptimWrapperConstructor`` will be used by default.
+
+    Args:
+        model (nn.Module): Model to be optimized.
+        cfg (dict): Config of optimizer wrapper, optimizer constructor and
+            optimizer.
+
+    Returns:
+        OptimWrapper: The built optimizer wrapper.
+    """
+    optim_wrapper_cfg = copy.deepcopy(cfg)
+    constructor_type = optim_wrapper_cfg.pop('constructor',
+                                             'DefaultOptimWrapperConstructor')
+    paramwise_cfg = optim_wrapper_cfg.pop('paramwise_cfg', None)
+
+    # Since the current generation of NPU(Ascend 910) only supports
+    # mixed precision training, here we turn on mixed precision by default
+    # on the NPU to make the training normal
+    if is_npu_available():
+        optim_wrapper_cfg['type'] = 'AmpOptimWrapper'
+
+    optim_wrapper_constructor = OPTIM_WRAPPER_CONSTRUCTORS.build(
+        dict(
+            type=constructor_type,
+            optim_wrapper_cfg=optim_wrapper_cfg,
+            paramwise_cfg=paramwise_cfg))
+    optim_wrapper = optim_wrapper_constructor(model)
+    return optim_wrapper

mmpl/engine/runner/__init__.py

@@ -0,0 +1,3 @@
+from .pl_runner import PLRunner
+
+__all__ = ['PLRunner']

mmpl/engine/runner/pl_runner.py

@@ -0,0 +1,941 @@
+import copy
+import logging
+import os
+import os.path as osp
+import pickle
+import platform
+import time
+import warnings
+from collections import OrderedDict
+from functools import partial
+from typing import Callable, Dict, List, Optional, Sequence, Union
+
+import torch
+import torch.nn as nn
+from lightning.pytorch.loggers import Logger
+from torch.nn.parallel.distributed import DistributedDataParallel
+from torch.optim import Optimizer
+from torch.utils.data import DataLoader
+
+import mmengine
+from mmengine.config import Config, ConfigDict
+from mmengine.dataset import worker_init_fn
+from mmengine.device import get_device
+from mmengine.dist import (broadcast, get_dist_info, get_rank, init_dist,
+                           is_distributed, master_only)
+from mmengine.evaluator import Evaluator
+from mmengine.fileio import FileClient, join_path
+from mmengine.hooks import Hook
+from mmengine.logging import MessageHub, MMLogger, print_log
+from mmengine.model import (MMDistributedDataParallel, convert_sync_batchnorm,
+                            is_model_wrapper, revert_sync_batchnorm)
+from mmengine.optim import (OptimWrapper, OptimWrapperDict, _ParamScheduler,
+                            build_optim_wrapper)
+from mmengine.registry import (DATA_SAMPLERS, DATASETS, EVALUATOR, FUNCTIONS,
+                               HOOKS, LOG_PROCESSORS, LOOPS, MODEL_WRAPPERS,
+                               OPTIM_WRAPPERS, PARAM_SCHEDULERS,
+                               RUNNERS, VISUALIZERS, DefaultScope)
+from mmengine.utils import digit_version, get_git_hash, is_seq_of
+from mmengine.utils.dl_utils import (TORCH_VERSION, collect_env,
+                                     set_multi_processing)
+from mmengine.visualization import Visualizer
+from mmengine.runner.base_loop import BaseLoop
+from mmengine.runner.checkpoint import (_load_checkpoint, _load_checkpoint_to_model,
+                         find_latest_checkpoint, get_state_dict,
+                         save_checkpoint, weights_to_cpu)
+from mmengine.runner.log_processor import LogProcessor
+from mmengine.runner.loops import EpochBasedTrainLoop, IterBasedTrainLoop, TestLoop, ValLoop
+from mmengine.runner.priority import Priority, get_priority
+from mmengine.runner.utils import set_random_seed
+
+ConfigType = Union[Dict, Config, ConfigDict]
+ParamSchedulerType = Union[List[_ParamScheduler], Dict[str, List[_ParamScheduler]]]
+OptimWrapperType = Union[OptimWrapper, OptimWrapperDict]
+
+from mmpl.registry import MODELS, LOGGERS
+import lightning.pytorch as pl
+from mmpl.models import build_pler
+
+
+@RUNNERS.register_module()
+class PLRunner:
+    def __init__(
+        self,
+        trainer_cfg: Dict,
+        model_cfg: Union[pl.LightningModule, Dict],
+        datamodule_cfg: Optional[Dict] = None,
+        cfg: Optional[ConfigType] = None
+    ):
+        self.trainer_cfg = copy.deepcopy(trainer_cfg)
+        self.model_cfg = copy.deepcopy(model_cfg)
+        self.datamodule_cfg = copy.deepcopy(datamodule_cfg)
+        mmengine.mkdir_or_exist(trainer_cfg['default_root_dir'])
+
+        timestamp = torch.tensor(time.time(), dtype=torch.float64)
+        # broadcast timestamp from 0 process to other processes
+        broadcast(timestamp)
+        self.timestamp = time.strftime('%Y%m%d_%H%M%S',
+                                       time.localtime(timestamp.item()))
+
+        if cfg is not None:
+            if isinstance(cfg, Config):
+                self.cfg = copy.deepcopy(cfg)
+            elif isinstance(cfg, dict):
+                self.cfg = Config(cfg)
+        else:
+            self.cfg = Config(dict())
+
+        compiled_model = trainer_cfg.pop('compiled_model', False)
+
+        # build logger
+        loggers = self.build_logger(
+            trainer_cfg.get('logger', False),
+            trainer_cfg.get('default_root_dir', f'{self.timestamp}')
+        )
+        trainer_cfg['logger'] = loggers
+
+        # build visualizer used for writing log or visualizing all kinds of data
+        self.visualizer = self.build_visualizer(
+            self.cfg.get('visualizer', None),
+            trainer_cfg.get('default_root_dir', f'{self.timestamp}')
+        )
+        if self.cfg:
+            self.visualizer.add_config(self.cfg)
+
+        # build callbacks
+        callbacks = self.build_hooks(
+            trainer_cfg.get('callbacks', None),
+        )
+        trainer_cfg['callbacks'] = callbacks
+
+        # build strategy
+        strategy = self.build_strategy(
+            trainer_cfg.get('strategy', 'auto'),
+        )
+        trainer_cfg['strategy'] = strategy
+
+        self.trainer = pl.Trainer(**trainer_cfg)
+        model_cfg.update({'config_cfg': copy.deepcopy(cfg).to_dict()})
+        model = self.build_model(model_cfg)
+        if cfg.get('load_from', None) is not None:
+            self.load_checkpoint(model, cfg['load_from'])
+        if compiled_model:
+            # default, reduce-overhead, and max-autotune.
+            self.model = torch.compile(model)
+        else:
+            self.model = model
+
+        # dump `cfg` to `work_dir`
+        self.dump_config()
+        # # Collect and log environment information.
+        # self._log_env(env_cfg)
+        # log hooks information
+        # self.logger.info(f'Hooks will be executed in the following '
+        #                  f'order:\n{self.get_hooks_info()}')
+
+    def build_visualizer(
+            self,
+            visualizer: Optional[Union[Visualizer,
+                                       Dict]] = None,
+            default_root_dir = 'tmp'
+    ) -> Visualizer:
+        """Build a global asscessable Visualizer.
+
+        Args:
+            visualizer (Visualizer or dict, optional): A Visualizer object
+                or a dict to build Visualizer object. If ``visualizer`` is a
+                Visualizer object, just returns itself. If not specified,
+                default config will be used to build Visualizer object.
+                Defaults to None.
+
+        Returns:
+            Visualizer: A Visualizer object build from ``visualizer``.
+        """
+        if visualizer is None:
+            visualizer = dict(
+                name=os.path.basename(default_root_dir),
+                vis_backends=[dict(type='LocalVisBackend')],
+                save_dir=default_root_dir+'/visualizer'
+            )
+            return Visualizer.get_instance(**visualizer)
+
+        if isinstance(visualizer, Visualizer):
+            return visualizer
+
+        if isinstance(visualizer, dict):
+            # ensure visualizer containing name key
+            visualizer.setdefault('name', os.path.basename(default_root_dir))
+            visualizer.setdefault('save_dir', default_root_dir+'/visualizer')
+            return VISUALIZERS.build(visualizer)
+        else:
+            raise TypeError(
+                'visualizer should be Visualizer object, a dict or None, '
+                f'but got {visualizer}')
+
+    def build_hooks(self, hooks: Union[Dict, List[Dict]] = None) -> List[Hook]:
+        """Build hooks from config.
+
+        Args:
+            hooks_cfg (dict): Config dict of hooks.
+
+        Returns:
+            list[Hook]: A list of hooks.
+        """
+        if hooks is not None:
+            if isinstance(hooks, dict):
+                hooks = [hooks]
+            tmp_hooks = []
+            for hook in hooks:
+                hook = HOOKS.build(hook)
+                tmp_hooks.append(hook)
+            hooks = tmp_hooks
+        return hooks
+
+    @classmethod
+    def from_cfg(cls, cfg: ConfigType) -> 'Runner':
+        cfg = copy.deepcopy(cfg)
+        runner = cls(
+            trainer_cfg=cfg.get('trainer_cfg'),
+            model_cfg=cfg['model_cfg'],
+            datamodule_cfg=cfg.get('datamodule_cfg'),
+            cfg=cfg
+        )
+
+        return runner
+
+    def build_logger(self, loggers: Union[Dict, List[Dict]] = None, default_root_dir='logger'):
+        if loggers is not None and loggers:
+            if isinstance(loggers, Dict):
+                loggers = [loggers]
+            tmp_loggers = []
+            for logger in loggers:
+                if logger.get('save_dir', None) is None:
+                    logger['save_dir'] = default_root_dir
+                mmengine.mkdir_or_exist(logger['save_dir'])
+                tmp_loggers.append(LOGGERS.build(logger))
+            loggers = tmp_loggers
+        return loggers
+
+    def build_strategy(self, strategy='auto'):
+        if isinstance(strategy, str):
+            return strategy
+        elif isinstance(strategy, dict):
+            if strategy.get('type', '') == 'FSDPStrategy':
+                from torch.distributed.fsdp import CPUOffload
+                from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
+                import functools
+                strategy.update(
+                    dict(
+                        # cpu_offload=CPUOffload(offload_params=True),
+                        auto_wrap_policy=functools.partial(
+                            size_based_auto_wrap_policy, min_num_params=int(5e7)
+                        )
+                    )
+                )
+            strategy = MODEL_WRAPPERS.build(strategy)
+            return strategy
+        return strategy
+
+    def build_model(self, model: Union[pl.LightningModule, Dict]) -> pl.LightningModule:
+        if isinstance(model, pl.LightningModule):
+            return model
+        elif isinstance(model, dict):
+            model = build_pler(model)
+            return model  # type: ignore
+        else:
+            raise TypeError('model should be a nn.Module object or dict, '
+                            f'but got {model}')
+
+    def _init_model_weights(self) -> None:
+        """Initialize the model weights if the model has
+        :meth:`init_weights`"""
+        if hasattr(self.model, 'module'):
+            model = self.model.module
+        else:
+            model = self.model
+        if hasattr(model, 'init_weights'):
+            model.init_weights()
+            # sync params and buffers
+            for name, params in model.state_dict().items():
+                broadcast(params)
+
+    def get_hooks_info(self) -> str:
+        # Get hooks info in each stage
+        stage_hook_map: Dict[str, list] = {stage: [] for stage in Hook.stages}
+        for hook in self.hooks:
+            try:
+                priority = Priority(hook.priority).name  # type: ignore
+            except ValueError:
+                priority = hook.priority  # type: ignore
+            classname = hook.__class__.__name__
+            hook_info = f'({priority:<12}) {classname:<35}'
+            for trigger_stage in hook.get_triggered_stages():
+                stage_hook_map[trigger_stage].append(hook_info)
+
+        stage_hook_infos = []
+        for stage in Hook.stages:
+            hook_infos = stage_hook_map[stage]
+            if len(hook_infos) > 0:
+                info = f'{stage}:\n'
+                info += '\n'.join(hook_infos)
+                info += '\n -------------------- '
+                stage_hook_infos.append(info)
+        return '\n'.join(stage_hook_infos)
+
+    def load_or_resume(self) -> None:
+        """load or resume checkpoint."""
+        if self._has_loaded:
+            return None
+
+        # decide to load from checkpoint or resume from checkpoint
+        resume_from = None
+        if self._resume and self._load_from is None:
+            # auto resume from the latest checkpoint
+            resume_from = find_latest_checkpoint(self.work_dir)
+            self.logger.info(
+                f'Auto resumed from the latest checkpoint {resume_from}.')
+        elif self._resume and self._load_from is not None:
+            # resume from the specified checkpoint
+            resume_from = self._load_from
+
+        if resume_from is not None:
+            self.resume(resume_from)
+            self._has_loaded = True
+        elif self._load_from is not None:
+            self.load_checkpoint(self._load_from)
+            self._has_loaded = True
+
+    @staticmethod
+    def build_datamodule(datamodule_cfg: Union[pl.LightningDataModule, Dict]):
+        if isinstance(datamodule_cfg, pl.LightningDataModule):
+            return datamodule_cfg
+        datamodule_cfg = copy.deepcopy(datamodule_cfg)
+        # build datamodule
+        datamodule = DATASETS.build(datamodule_cfg)
+        return datamodule
+
+    def run(self, status, *args, **kwargs):
+        assert status in ['fit', 'test', 'predict', 'validate']
+        trainer_func = self.trainer.__getattribute__(status)
+        self.datamodule = self.build_datamodule(self.datamodule_cfg)
+        return trainer_func(model=self.model, datamodule=self.datamodule, *args, **kwargs)
+
+        #
+        # if is_model_wrapper(self.model):
+        #     ori_model = self.model.module
+        # else:
+        #     ori_model = self.model
+        # assert hasattr(ori_model, 'train_step'), (
+        #     'If you want to train your model, please make sure your model '
+        #     'has implemented `train_step`.')
+        #
+        # if self._val_loop is not None:
+        #     assert hasattr(ori_model, 'val_step'), (
+        #         'If you want to validate your model, please make sure your '
+        #         'model has implemented `val_step`.')
+        #
+        # if self._train_loop is None:
+        #     raise RuntimeError(
+        #         '`self._train_loop` should not be None when calling train '
+        #         'method. Please provide `train_dataloader`, `train_cfg`, '
+        #         '`optimizer` and `param_scheduler` arguments when '
+        #         'initializing runner.')
+        #
+        # self._train_loop = self.build_train_loop(
+        #     self._train_loop)  # type: ignore
+        #
+        # # `build_optimizer` should be called before `build_param_scheduler`
+        # #  because the latter depends on the former
+        # self.optim_wrapper = self.build_optim_wrapper(self.optim_wrapper)
+        # # Automatically scaling lr by linear scaling rule
+        # self.scale_lr(self.optim_wrapper, self.auto_scale_lr)
+        #
+        # if self.param_schedulers is not None:
+        #     self.param_schedulers = self.build_param_scheduler(  # type: ignore
+        #         self.param_schedulers)  # type: ignore
+        #
+        # if self._val_loop is not None:
+        #     self._val_loop = self.build_val_loop(
+        #         self._val_loop)  # type: ignore
+        # # TODO: add a contextmanager to avoid calling `before_run` many times
+        # self.call_hook('before_run')
+        #
+        # # initialize the model weights
+        # self._init_model_weights()
+        # # make sure checkpoint-related hooks are triggered after `before_run`
+        # self.load_or_resume()
+        #
+        # # Initiate inner count of `optim_wrapper`.
+        # self.optim_wrapper.initialize_count_status(
+        #     self.model,
+        #     self._train_loop.iter,  # type: ignore
+        #     self._train_loop.max_iters)  # type: ignore
+        #
+        # # Maybe compile the model according to options in self.cfg.compile
+        # # This must be called **AFTER** model has been wrapped.
+        # self._maybe_compile('train_step')
+        #
+        # model = self.train_loop.run()  # type: ignore
+        # self.call_hook('after_run')
+        # return model
+
+
+
+    def register_hook(
+            self,
+            hook: Union[Hook, Dict],
+            priority: Optional[Union[str, int, Priority]] = None) -> None:
+        """Register a hook into the hook list.
+
+        The hook will be inserted into a priority queue, with the specified
+        priority (See :class:`Priority` for details of priorities).
+        For hooks with the same priority, they will be triggered in the same
+        order as they are registered.
+
+        Priority of hook will be decided with the following priority:
+
+        - ``priority`` argument. If ``priority`` is given, it will be priority
+          of hook.
+        - If ``hook`` argument is a dict and ``priority`` in it, the priority
+          will be the value of ``hook['priority']``.
+        - If ``hook`` argument is a dict but ``priority`` not in it or ``hook``
+          is an instance of ``hook``, the priority will be ``hook.priority``.
+
+        Args:
+            hook (:obj:`Hook` or dict): The hook to be registered.
+            priority (int or str or :obj:`Priority`, optional): Hook priority.
+                Lower value means higher priority.
+        """
+        if not isinstance(hook, (Hook, dict)):
+            raise TypeError(
+                f'hook should be an instance of Hook or dict, but got {hook}')
+
+        _priority = None
+        if isinstance(hook, dict):
+            if 'priority' in hook:
+                _priority = hook.pop('priority')
+
+            hook_obj = HOOKS.build(hook)
+        else:
+            hook_obj = hook
+
+        if priority is not None:
+            hook_obj.priority = priority
+        elif _priority is not None:
+            hook_obj.priority = _priority
+
+        inserted = False
+        for i in range(len(self._hooks) - 1, -1, -1):
+            if get_priority(hook_obj.priority) >= get_priority(
+                    self._hooks[i].priority):
+                self._hooks.insert(i + 1, hook_obj)
+                inserted = True
+                break
+        if not inserted:
+            self._hooks.insert(0, hook_obj)
+
+    def register_default_hooks(
+            self,
+            hooks: Optional[Dict[str, Union[Hook, Dict]]] = None) -> None:
+        """Register default hooks into hook list.
+
+        ``hooks`` will be registered into runner to execute some default
+        actions like updating model parameters or saving checkpoints.
+
+        Default hooks and their priorities:
+
+        +----------------------+-------------------------+
+        | Hooks                | Priority                |
+        +======================+=========================+
+        | RuntimeInfoHook      | VERY_HIGH (10)          |
+        +----------------------+-------------------------+
+        | IterTimerHook        | NORMAL (50)             |
+        +----------------------+-------------------------+
+        | DistSamplerSeedHook  | NORMAL (50)             |
+        +----------------------+-------------------------+
+        | LoggerHook           | BELOW_NORMAL (60)       |
+        +----------------------+-------------------------+
+        | ParamSchedulerHook   | LOW (70)                |
+        +----------------------+-------------------------+
+        | CheckpointHook       | VERY_LOW (90)           |
+        +----------------------+-------------------------+
+
+        If ``hooks`` is None, above hooks will be registered by
+        default::
+
+            default_hooks = dict(
+                runtime_info=dict(type='RuntimeInfoHook'),
+                timer=dict(type='IterTimerHook'),
+                sampler_seed=dict(type='DistSamplerSeedHook'),
+                logger=dict(type='LoggerHook'),
+                param_scheduler=dict(type='ParamSchedulerHook'),
+                checkpoint=dict(type='CheckpointHook', interval=1),
+            )
+
+        If not None, ``hooks`` will be merged into ``default_hooks``.
+        If there are None value in default_hooks, the corresponding item will
+        be popped from ``default_hooks``::
+
+            hooks = dict(timer=None)
+
+        The final registered default hooks will be :obj:`RuntimeInfoHook`,
+        :obj:`DistSamplerSeedHook`, :obj:`LoggerHook`,
+        :obj:`ParamSchedulerHook` and :obj:`CheckpointHook`.
+
+        Args:
+            hooks (dict[str, Hook or dict], optional): Default hooks or configs
+                to be registered.
+        """
+        default_hooks: dict = dict(
+            runtime_info=dict(type='RuntimeInfoHook'),
+            timer=dict(type='IterTimerHook'),
+            sampler_seed=dict(type='DistSamplerSeedHook'),
+            logger=dict(type='LoggerHook'),
+            param_scheduler=dict(type='ParamSchedulerHook'),
+            checkpoint=dict(type='CheckpointHook', interval=1),
+        )
+        if hooks is not None:
+            for name, hook in hooks.items():
+                if name in default_hooks and hook is None:
+                    # remove hook from _default_hooks
+                    default_hooks.pop(name)
+                else:
+                    assert hook is not None
+                    default_hooks[name] = hook
+
+        for hook in default_hooks.values():
+            self.register_hook(hook)
+
+    def register_custom_hooks(self, hooks: List[Union[Hook, Dict]]) -> None:
+        """Register custom hooks into hook list.
+
+        Args:
+            hooks (list[Hook | dict]): List of hooks or configs to be
+                registered.
+        """
+        for hook in hooks:
+            self.register_hook(hook)
+
+    def register_hooks(
+            self,
+            default_hooks: Optional[Dict[str, Union[Hook, Dict]]] = None,
+            custom_hooks: Optional[List[Union[Hook, Dict]]] = None) -> None:
+        """Register default hooks and custom hooks into hook list.
+
+        Args:
+            default_hooks (dict[str, dict] or dict[str, Hook], optional): Hooks
+                to execute default actions like updating model parameters and
+                saving checkpoints.  Defaults to None.
+            custom_hooks (list[dict] or list[Hook], optional): Hooks to execute
+                custom actions like visualizing images processed by pipeline.
+                Defaults to None.
+        """
+        self.register_default_hooks(default_hooks)
+
+        if custom_hooks is not None:
+            self.register_custom_hooks(custom_hooks)
+
+    def resume(self,
+               filename: str,
+               resume_optimizer: bool = True,
+               resume_param_scheduler: bool = True,
+               map_location: Union[str, Callable] = 'default') -> None:
+        """Resume model from checkpoint.
+
+        Args:
+            filename (str): Accept local filepath, URL, ``torchvision://xxx``,
+                ``open-mmlab://xxx``.
+            resume_optimizer (bool): Whether to resume optimizer state.
+                Defaults to True.
+            resume_param_scheduler (bool): Whether to resume param scheduler
+                state. Defaults to True.
+            map_location (str or callable):A string or a callable function to
+                specifying how to remap storage locations.
+                Defaults to 'default'.
+        """
+        if map_location == 'default':
+            device = get_device()
+            checkpoint = self.load_checkpoint(filename, map_location=device)
+        else:
+            checkpoint = self.load_checkpoint(
+                filename, map_location=map_location)
+
+        self.train_loop._epoch = checkpoint['meta']['epoch']
+        self.train_loop._iter = checkpoint['meta']['iter']
+
+        # check whether the number of GPU used for current experiment
+        # is consistent with resuming from checkpoint
+        if 'config' in checkpoint['meta']:
+            config = mmengine.Config.fromstring(
+                checkpoint['meta']['config'], file_format='.py')
+            previous_gpu_ids = config.get('gpu_ids', None)
+            if (previous_gpu_ids is not None and len(previous_gpu_ids) > 0
+                    and len(previous_gpu_ids) != self._world_size):
+                # TODO, should we modify the iteration?
+                self.logger.info(
+                    'Number of GPU used for current experiment is not '
+                    'consistent with resuming from checkpoint')
+                if (self.auto_scale_lr is None
+                        or not self.auto_scale_lr.get('enable', False)):
+                    raise RuntimeError(
+                        'Cannot automatically rescale lr in resuming. Please '
+                        'make sure the number of GPU is consistent with the '
+                        'previous training state resuming from the checkpoint '
+                        'or set `enable` in `auto_scale_lr to False.')
+
+        # resume random seed
+        resumed_seed = checkpoint['meta'].get('seed', None)
+        current_seed = self._randomness_cfg.get('seed')
+        if resumed_seed is not None and resumed_seed != current_seed:
+            if current_seed is not None:
+                print_log(
+                    f'The value of random seed in the '
+                    f'checkpoint "{resumed_seed}" is '
+                    f'different from the value in '
+                    f'`randomness` config "{current_seed}"',
+                    logger='current',
+                    level=logging.WARNING)
+            self._randomness_cfg.update(seed=resumed_seed)
+            self.set_randomness(**self._randomness_cfg)
+
+        resumed_dataset_meta = checkpoint['meta'].get('dataset_meta', None)
+        dataset_meta = getattr(self.train_dataloader.dataset, 'metainfo', None)
+
+        # `resumed_dataset_meta` and `dataset_meta` could be object like
+        # np.ndarray, which cannot be directly judged as equal or not,
+        # therefore we just compared their dumped results.
+        if pickle.dumps(resumed_dataset_meta) != pickle.dumps(dataset_meta):
+            print_log(
+                'The dataset metainfo from the resumed checkpoint is '
+                'different from the current training dataset, please '
+                'check the correctness of the checkpoint or the training '
+                'dataset.',
+                logger='current',
+                level=logging.WARNING)
+
+        self.message_hub.load_state_dict(checkpoint['message_hub'])
+
+        # resume optimizer
+        if 'optimizer' in checkpoint and resume_optimizer:
+            self.optim_wrapper = self.build_optim_wrapper(self.optim_wrapper)
+            self.optim_wrapper.load_state_dict(  # type: ignore
+                checkpoint['optimizer'])
+
+        # resume param scheduler
+        if resume_param_scheduler and self.param_schedulers is None:
+            print_log(
+                '`resume_param_scheduler` is True but `self.param_schedulers` '
+                'is None, so skip resuming parameter schedulers',
+                logger='current',
+                level=logging.WARNING)
+            resume_param_scheduler = False
+        if 'param_schedulers' in checkpoint and resume_param_scheduler:
+            self.param_schedulers = self.build_param_scheduler(  # type: ignore
+                self.param_schedulers)  # type: ignore
+            if isinstance(self.param_schedulers, dict):
+                for name, schedulers in self.param_schedulers.items():
+                    for scheduler, ckpt_scheduler in zip(
+                            schedulers, checkpoint['param_schedulers'][name]):
+                        scheduler.load_state_dict(ckpt_scheduler)
+            else:
+                for scheduler, ckpt_scheduler in zip(
+                        self.param_schedulers,  # type: ignore
+                        checkpoint['param_schedulers']):
+                    scheduler.load_state_dict(ckpt_scheduler)
+
+        self._has_loaded = True
+
+        self.logger.info(f'resumed epoch: {self.epoch}, iter: {self.iter}')
+
+    # def load_checkpoint(self,
+    #                     filename: str,
+    #                     model,
+    #                     map_location: Union[str, Callable] = 'cpu',
+    #                     strict: bool = False,
+    #                     revise_keys: list = [(r'^module.', '')]):
+    #     """Load checkpoint from given ``filename``.
+    #
+    #     Args:
+    #         filename (str): Accept local filepath, URL, ``torchvision://xxx``,
+    #             ``open-mmlab://xxx``.
+    #         map_location (str or callable): A string or a callable function to
+    #             specifying how to remap storage locations.
+    #             Defaults to 'cpu'.
+    #         strict (bool): strict (bool): Whether to allow different params for
+    #             the model and checkpoint.
+    #         revise_keys (list): A list of customized keywords to modify the
+    #             state_dict in checkpoint. Each item is a (pattern, replacement)
+    #             pair of the regular expression operations. Defaults to strip
+    #             the prefix 'module.' by [(r'^module\\.', '')].
+    #     """
+    #     checkpoint = _load_checkpoint(filename, map_location=map_location)
+    #
+    #     if is_model_wrapper(model):
+    #         model = model.module
+    #     else:
+    #         model = model
+    #
+    #     checkpoint = _load_checkpoint_to_model(
+    #         model, checkpoint, strict, revise_keys=revise_keys)
+    #
+    #     print(f'Load checkpoint from {filename}')
+    #
+    #     return checkpoint
+    def load_checkpoint(self, model, file):
+
+        if isinstance(file, str):
+            file_path = file
+            state_dict = torch.load(file_path, map_location='cpu')['state_dict']
+        elif isinstance(file, dict):
+            file_path = file['file_path']
+            state_dict = torch.load(file_path, map_location='cpu')['state_dict']
+            for delete_key in file['delete_keys']:
+                del state_dict[delete_key]
+        else:
+            raise TypeError('file must be str or dict')
+        missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+        print('load from:', file_path)
+        print('load model missing_keys:', missing_keys)
+        print('load model unexpected_keys:', unexpected_keys)
+
+    @master_only
+    def save_checkpoint(
+        self,
+        out_dir: str,
+        filename: str,
+        file_client_args: Optional[dict] = None,
+        save_optimizer: bool = True,
+        save_param_scheduler: bool = True,
+        meta: dict = None,
+        by_epoch: bool = True,
+        backend_args: Optional[dict] = None,
+    ):
+        """Save checkpoints.
+
+        ``CheckpointHook`` invokes this method to save checkpoints
+        periodically.
+
+        Args:
+            out_dir (str): The directory that checkpoints are saved.
+            filename (str): The checkpoint filename.
+            file_client_args (dict, optional): Arguments to instantiate a
+                FileClient. See :class:`mmengine.fileio.FileClient` for
+                details. Defaults to None. It will be deprecated in future.
+                Please use `backend_args` instead.
+            save_optimizer (bool): Whether to save the optimizer to
+                the checkpoint. Defaults to True.
+            save_param_scheduler (bool): Whether to save the param_scheduler
+                to the checkpoint. Defaults to True.
+            meta (dict, optional): The meta information to be saved in the
+                checkpoint. Defaults to None.
+            by_epoch (bool): Whether the scheduled momentum is updated by
+                epochs. Defaults to True.
+            backend_args (dict, optional): Arguments to instantiate the
+                prefix of uri corresponding backend. Defaults to None.
+                New in v0.2.0.
+        """
+        if meta is None:
+            meta = {}
+        elif not isinstance(meta, dict):
+            raise TypeError(
+                f'meta should be a dict or None, but got {type(meta)}')
+
+        if by_epoch:
+            # self.epoch increments 1 after
+            # `self.call_hook('after_train_epoch)` but `save_checkpoint` is
+            # called by `after_train_epoch`` method of `CheckpointHook` so
+            # `epoch` should be `self.epoch + 1`
+            meta.update(epoch=self.epoch + 1, iter=self.iter)
+        else:
+            meta.update(epoch=self.epoch, iter=self.iter + 1)
+
+        if file_client_args is not None:
+            warnings.warn(
+                '"file_client_args" will be deprecated in future. '
+                'Please use "backend_args" instead', DeprecationWarning)
+            if backend_args is not None:
+                raise ValueError(
+                    '"file_client_args" and "backend_args" cannot be set at '
+                    'the same time.')
+
+            file_client = FileClient.infer_client(file_client_args, out_dir)
+            filepath = file_client.join_path(out_dir, filename)
+        else:
+            filepath = join_path(  # type: ignore
+                out_dir, filename, backend_args=backend_args)
+
+        meta.update(
+            cfg=self.cfg.pretty_text,
+            seed=self.seed,
+            experiment_name=self.experiment_name,
+            time=time.strftime('%Y%m%d_%H%M%S', time.localtime()),
+            mmengine_version=mmengine.__version__ + get_git_hash())
+
+        if hasattr(self.train_dataloader.dataset, 'metainfo'):
+            meta.update(dataset_meta=self.train_dataloader.dataset.metainfo)
+
+        if is_model_wrapper(self.model):
+            model = self.model.module
+        else:
+            model = self.model
+
+        checkpoint = {
+            'meta': meta,
+            'state_dict': weights_to_cpu(get_state_dict(model)),
+            'message_hub': self.message_hub.state_dict()
+        }
+        # save optimizer state dict to checkpoint
+        if save_optimizer:
+            if isinstance(self.optim_wrapper, OptimWrapper):
+                checkpoint['optimizer'] = self.optim_wrapper.state_dict()
+            else:
+                raise TypeError(
+                    'self.optim_wrapper should be an `OptimWrapper` '
+                    'or `OptimWrapperDict` instance, but got '
+                    f'{self.optim_wrapper}')
+
+        # save param scheduler state dict
+        if save_param_scheduler and self.param_schedulers is None:
+            print_log(
+                '`save_param_scheduler` is True but `self.param_schedulers` '
+                'is None, so skip saving parameter schedulers',
+                logger='current',
+                level=logging.WARNING)
+            save_param_scheduler = False
+        if save_param_scheduler:
+            if isinstance(self.param_schedulers, dict):
+                checkpoint['param_schedulers'] = dict()
+                for name, schedulers in self.param_schedulers.items():
+                    checkpoint['param_schedulers'][name] = []
+                    for scheduler in schedulers:
+                        state_dict = scheduler.state_dict()
+                        checkpoint['param_schedulers'][name].append(state_dict)
+            else:
+                checkpoint['param_schedulers'] = []
+                for scheduler in self.param_schedulers:  # type: ignore
+                    state_dict = scheduler.state_dict()  # type: ignore
+                    checkpoint['param_schedulers'].append(state_dict)
+
+        self.call_hook('before_save_checkpoint', checkpoint=checkpoint)
+        save_checkpoint(checkpoint, filepath)
+
+    @master_only
+    def dump_config(self) -> None:
+        version = ''
+        if len(self.trainer.loggers) > 0:
+            version = self.trainer.loggers[0].version
+            version = version if isinstance(version, str) else f"version_{version}"
+        if version == '':
+            # if no loggers, use default_root_dir
+            version = 'version'
+
+        """Dump config to `work_dir`."""
+        if self.cfg.filename is not None:
+            filename = osp.basename(self.cfg.filename)
+        else:
+            filename = f'{self.timestamp}.py'
+        path = f'{self.trainer.default_root_dir}/{version}_{filename}'
+
+        self.cfg.dump(path)
+
+    def _check_scheduler_cfg(
+            self, param_scheduler: Optional[Union[dict, list,
+                                                  _ParamScheduler]]) -> None:
+        """Parse `param_scheduler` to a list of parameter schedulers, or a
+        `dict` of which each value is a list of parameter schedulers.
+
+        If only one optimizer is used, the parsed config should be a
+        list of parameter scheduler configs or instances. If multiple
+        optimizers are used, the parsed config should be `dict`.
+        Its key should be consistent with the optimizer `dict` and its value
+        should be a list of parameter scheduler configs or instances. See
+        :meth:`build_param_scheduler` for more details.
+
+        Examples:
+            >>> # valid scheduler:
+            >>> # empty scheduler
+            >>> scheduler = None
+            >>> # Single scheduler
+            >>> scheduler = dict(type='MultiStepLR', milestones=[1, 2])
+            >>> # Single list schedulers
+            >>> scheduler = [dict(type='MultiStepLR', milestones=[1, 2]),
+            >>>              dict(type='MultiStepLR', milestones=[2, 3])]
+            >>> # `dict` of schedulers
+            >>> scheduler = dict(linear1=dict(type='MultiStepLR', milestones=[1, 2]),
+            >>>                  linear2=dict(type='MultiStepLR', milestones=[1, 2]))
+            >>> # `dict` of `list` of schedulers
+            >>> scheduler = dict(linear1=[dict(type='MultiStepLR', milestones=[1, 2])],
+            >>>                  linear2=[dict(type='MultiStepLR', milestones=[1, 2])])
+            >>> # Single built scheduler
+            >>> from mmengine.optim import MultiStepLR
+            >>> scheduler = MultiStepLR(milestones=[1, 2], optimizer=optimizer)
+            >>> # Single built list schedulers
+            >>> scheduler = [MultiStepLR(milestones=[1, 2], optimizer=optimizer)]
+            >>> # dict of built scheduler
+            >>> scheduler = dict(linear1=MultiStepLR(milestones=[1, 2], optimizer=optimizer),
+            >>>                  linear2=MultiStepLR(milestones=[1, 2], optimizer=optimizer))
+            >>> # dict of built list schedulers
+            >>> scheduler = dict(linear1=[MultiStepLR(milestones=[1, 2], optimizer=optimizer)],
+            >>>                  linear2=[MultiStepLR(milestones=[1, 2], optimizer=optimizer)])
+
+        Args:
+            param_scheduler (dict or list): The original parameter scheduler.
+        """  # noqa: E501
+        param_schedulers: Union[dict, list, _ParamScheduler]
+        if param_scheduler is None:
+            return
+        if isinstance(param_scheduler, _ParamScheduler):
+            return
+        if is_seq_of(param_scheduler, _ParamScheduler):
+            return
+
+        if is_seq_of(param_scheduler, dict):
+            for _param_scheduler in param_scheduler:
+                assert 'type' in _param_scheduler, (
+                    'Each parameter scheduler should contain the key type, '
+                    f'but got {_param_scheduler}')
+        elif isinstance(param_scheduler, dict):
+            if 'type' not in param_scheduler:
+                for key, _param_scheduler in param_scheduler.items():
+                    assert isinstance(
+                        _param_scheduler,
+                        (dict, tuple, list, _ParamScheduler)), (
+                            'Each value of `param_scheduler` should be a '
+                            f'dict or a list, but got {_param_scheduler} with '
+                            f'type {type(_ParamScheduler)}')
+
+        else:
+            raise TypeError(
+                '`param_scheduler` should be a `_ParamScheduler`, `dict`, '
+                f'list or a tuple, but got {type(param_scheduler)}. If '
+                '`param_scheduler` is a list of dict, it means a list of '
+                'scheduler configs for single optimizer. If it is a dict and '
+                'contains key `type`, it means a scheduler config for a '
+                'single optimizer. If it does not contain key `type`, it '
+                'means multiple lists of schedulers for multiple optimizers.')
+
+    def _log_env(self, env_cfg: dict) -> None:
+        """Logging environment information of the current task.
+
+        Args:
+            env_cfg (dict): The environment config of the runner.
+        """
+        # Collect and log environment information.
+        env = collect_env()
+        runtime_env = OrderedDict()
+        runtime_env.update(env_cfg)
+        runtime_env.update(self._randomness_cfg)
+        runtime_env['Distributed launcher'] = self._launcher
+        runtime_env['Distributed training'] = self._distributed
+        runtime_env['GPU number'] = self._world_size
+
+        env_info = '\n    ' + '\n    '.join(f'{k}: {v}'
+                                            for k, v in env.items())
+        runtime_env_info = '\n    ' + '\n    '.join(
+            f'{k}: {v}' for k, v in runtime_env.items())
+        dash_line = '-' * 60
+        self.logger.info('\n' + dash_line + '\nSystem environment:' +
+                         env_info + '\n'
+                         '\nRuntime environment:' + runtime_env_info + '\n' +
+                         dash_line + '\n')
+        self.logger.info(f'Config:\n{self.cfg.pretty_text}')

mmpl/engine/strategies/__init__.py

@@ -0,0 +1 @@
+from .builder import PL_MODEL_WRAPPERS