Remove runner, parallel, engine and device (#2216)

* Remove runner, parallel, engine and device

* fix format

* remove outdated docs

MANIFEST.in

 include requirements/runtime.txt
-include mmcv/model_zoo/open_mmlab.json mmcv/model_zoo/deprecated.json mmcv/model_zoo/mmcls.json mmcv/model_zoo/torchvision_0.12.json
 include mmcv/ops/csrc/common/cuda/*.cuh mmcv/ops/csrc/common/cuda/*.hpp mmcv/ops/csrc/common/*.hpp
 include mmcv/ops/csrc/pytorch/*.cpp mmcv/ops/csrc/pytorch/cuda/*.cu mmcv/ops/csrc/pytorch/cuda/*.cpp mmcv/ops/csrc/pytorch/cpu/*.cpp
 include mmcv/ops/csrc/parrots/*.h mmcv/ops/csrc/parrots/*.cpp

docs/en/api.rst

@@ -28,16 +28,6 @@ cnn
 .. automodule:: mmcv.cnn
     :members:
 
-runner
-------
-.. automodule:: mmcv.runner
-    :members:
-
-engine
-------
-.. automodule:: mmcv.engine
-    :members:
-
 ops
 ------
 .. automodule:: mmcv.ops

docs/en/index.rst

@@ -17,7 +17,6 @@ You can switch between Chinese and English documents in the lower-left corner of
 
    understand_mmcv/config.md
    understand_mmcv/registry.md
-   understand_mmcv/runner.md
    understand_mmcv/data_process.md
    understand_mmcv/visualization.md
    understand_mmcv/cnn.md

docs/en/understand_mmcv/runner.md

-## Runner
-
-The runner class is designed to manage the training. It eases the training process with less code demanded from users while staying flexible and configurable. The main features are as listed:
-
-- Support `EpochBasedRunner` and `IterBasedRunner` for different scenarios. Implementing customized runners is also allowed to meet customized needs.
-- Support customized workflow to allow switching between different modes while training. Currently, supported modes are train and val.
-- Enable extensibility through various hooks, including hooks defined in MMCV and customized ones.
-
-### EpochBasedRunner
-
-As its name indicates, workflow in `EpochBasedRunner` should be set based on epochs. For example, \[('train', 2), ('val', 1)\] means running 2 epochs for training and 1 epoch for validation, iteratively. And each epoch may contain multiple iterations. Currently, MMDetection uses `EpochBasedRunner` by default.
-
-Let's take a look at its core logic:
-
-```python
-# the condition to stop training
-while curr_epoch < max_epochs:
-    # traverse the workflow.
-    # e.g. workflow = [('train', 2), ('val', 1)]
-    for i, flow in enumerate(workflow):
-        # mode(e.g. train) determines which function to run
-        mode, epochs = flow
-        # epoch_runner will be either self.train() or self.val()
-        epoch_runner = getattr(self, mode)
-        # execute the corresponding function
-        for _ in range(epochs):
-            epoch_runner(data_loaders[i], **kwargs)
-```
-
-Currently, we support 2 modes: train and val. Let's take a train function for example and have a look at its core logic:
-
-```python
-# Currently, epoch_runner could be either train or val
-def train(self, data_loader, **kwargs):
-    # traverse the dataset and get batch data for 1 epoch
-    for i, data_batch in enumerate(data_loader):
-        # it will execute all before_train_iter function in the hooks registered. You may want to watch out for the order.
-        self.call_hook('before_train_iter')
-        # set train_mode as False in val function
-        self.run_iter(data_batch, train_mode=True, **kwargs)
-        self.call_hook('after_train_iter')
-   self.call_hook('after_train_epoch')
-```
-
-### IterBasedRunner
-
-Different from `EpochBasedRunner`, workflow in `IterBasedRunner` should be set based on iterations. For example, \[('train', 2), ('val', 1)\] means running 2 iters for training and 1 iter for validation, iteratively. Currently, MMSegmentation uses `IterBasedRunner` by default.
-
-Let's take a look at its core logic:
-
-```python
-# Although we set workflow by iters here, we might also need info on the epochs in some using cases. That can be provided by IterLoader.
-iter_loaders = [IterLoader(x) for x in data_loaders]
-# the condition to stop training
-while curr_iter < max_iters:
-    # traverse the workflow.
-    # e.g. workflow = [('train', 2), ('val', 1)]
-    for i, flow in enumerate(workflow):
-        # mode(e.g. train) determines which function to run
-        mode, iters = flow
-        # iter_runner will be either self.train() or self.val()
-        iter_runner = getattr(self, mode)
-        # execute the corresponding function
-        for _ in range(iters):
-            iter_runner(iter_loaders[i], **kwargs)
-```
-
-Currently, we support 2 modes: train and val. Let's take a val function for example and have a look at its core logic:
-
-```python
-# Currently, iter_runner could be either train or val
-def val(self, data_loader, **kwargs):
-    # get batch data for 1 iter
-    data_batch = next(data_loader)
-    # it will execute all before_val_iter function in the hooks registered. You may want to watch out for the order.
-    self.call_hook('before_val_iter')
-    outputs = self.model.val_step(data_batch, self.optimizer, **kwargs)
-    self.outputs = outputs
-    self.call_hook('after_val_iter')
-```
-
-Other than the basic functionalities explained above, `EpochBasedRunner` and `IterBasedRunner` provide methods such as `resume`, `save_checkpoint` and `register_hook`. In case you are not familiar with the term Hook mentioned earlier, we will also provide a tutorial about it.(coming soon...) Essentially, a hook is functionality to alter or augment the code behaviors through predefined api. It allows users to have their own code called under certain circumstances. It makes code extensible in a non-intrusive manner.
-
-### A Simple Example
-
-We will walk you through the usage of runner with a classification task. The following code only contains essential steps for demonstration purposes. The following steps are necessary for any training tasks.
-
-**(1) Initialize dataloader, model, optimizer, etc.**
-
-```python
-# initialize model
-model=...
-# initialize optimizer, typically, we set: cfg.optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0001)
-optimizer = build_optimizer(model, cfg.optimizer)
-# initialize the dataloader corresponding to the workflow(train/val)
-data_loaders = [
-        build_dataloader(
-            ds,
-            cfg.data.samples_per_gpu,
-            cfg.data.workers_per_gpu,
-            ...) for ds in dataset
-    ]
-```
-
-**(2) Initialize runner**
-
-```python
-runner = build_runner(
-    # cfg.runner is typically set as:
-    # runner = dict(type='EpochBasedRunner', max_epochs=200)
-    cfg.runner,
-    default_args=dict(
-        model=model,
-        batch_processor=None,
-        optimizer=optimizer,
-        logger=logger))
-```
-
-**(3) Register training hooks and customized hooks.**
-
-```python
-# register default hooks necessary for training
-runner.register_training_hooks(
-    # configs of learning rate, it is typically set as:
-    # lr_config = dict(policy='step', step=[100, 150])
-    cfg.lr_config,
-    # configuration of optimizer, e.g. grad_clip
-    optimizer_config,
-    # configuration of saving checkpoints, it is typically set as:
-    # checkpoint_config = dict(interval=1), saving checkpoints every epochs
-    cfg.checkpoint_config,
-    # configuration of logs
-    cfg.log_config,
-    ...)
-
-# register customized hooks
-# say we want to enable ema, then we could set custom_hooks=[dict(type='EMAHook')]
-if cfg.get('custom_hooks', None):
-    custom_hooks = cfg.custom_hooks
-    for hook_cfg in cfg.custom_hooks:
-        hook_cfg = hook_cfg.copy()
-        priority = hook_cfg.pop('priority', 'NORMAL')
-        hook = build_from_cfg(hook_cfg, HOOKS)
-        runner.register_hook(hook, priority=priority)
-```
-
-Then, we can use `resume` or `load_checkpoint` to load existing weights.
-
-**(4) Start training**
-
-```python
-# workflow is typically set as: workflow = [('train', 1)]
-# here the training begins.
-runner.run(data_loaders, cfg.workflow)
-```
-
-Let's take `EpochBasedRunner` for example and go a little bit into details about setting workflow:
-
-- Say we only want to put train in the workflow, then we can set: workflow = \[('train', 1)\]. The runner will only execute train iteratively in this case.
-- Say we want to put both train and val in the workflow, then we can set: workflow = \[('train', 3), ('val',1)\]. The runner will first execute train for 3 epochs and then switch to val mode and execute val for 1 epoch. The workflow will be repeated until the current epoch hit the max_epochs.
-- Workflow is highly flexible. Therefore, you can set workflow = \[('val', 1), ('train',1)\] if you would like the runner to validate first and train after.
-
-The code we demonstrated above is already in `train.py` in MM repositories. Simply modify the corresponding keys in the configuration files and the script will execute the expected workflow automatically.

docs/zh_cn/api.rst

@@ -28,16 +28,6 @@ cnn
 .. automodule:: mmcv.cnn
     :members:
 
-runner
-------
-.. automodule:: mmcv.runner
-    :members:
-
-engine
-------
-.. automodule:: mmcv.engine
-    :members:
-
 ops
 ------
 .. automodule:: mmcv.ops

docs/zh_cn/index.rst

@@ -17,7 +17,6 @@
 
    understand_mmcv/config.md
    understand_mmcv/registry.md
-   understand_mmcv/runner.md
    understand_mmcv/data_process.md
    understand_mmcv/data_transform.md
    understand_mmcv/visualization.md

docs/zh_cn/understand_mmcv/runner.md

-## 执行器
-
-执行器模块负责模型训练过程调度，主要目的是让用户使用更少的代码以及灵活可配置方式开启训练。其具备如下核心特性:
-
-- 支持以 `EpochBasedRunner` 和 `IterBasedRunner` 为单位的迭代模式以满足不同场景
-- 支持定制工作流以满足训练过程中各状态自由切换，目前支持训练和验证两个工作流。工作流可以简单理解为一个完成的训练和验证迭代过程。
-- 配合各类默认和自定义 Hook，对外提供了灵活扩展能力
-
-### EpochBasedRunner
-
-顾名思义，`EpochBasedRunner` 是指以 epoch 为周期的工作流，例如设置 workflow = \[('train', 2), ('val', 1)\] 表示循环迭代地训练 2 个 epoch，然后验证 1 个 epoch。MMDetection 目标检测框架默认采用的是 `EpochBasedRunner`。
-
-其抽象逻辑如下所示：
-
-```python
-# 训练终止条件
-while curr_epoch < max_epochs:
-    # 遍历用户设置的工作流，例如 workflow = [('train', 2)，('val', 1)]
-    for i, flow in enumerate(workflow):
-        # mode 是工作流函数，例如 train, epochs 是迭代次数
-        mode, epochs = flow
-        # 要么调用 self.train()，要么调用 self.val()
-        epoch_runner = getattr(self, mode)
-        # 运行对应工作流函数
-        for _ in range(epochs):
-            epoch_runner(data_loaders[i], **kwargs)
-```
-
-目前支持训练和验证两个工作流，以训练函数为例，其抽象逻辑是：
-
-```python
-# epoch_runner 目前可以是 train 或者 val
-def train(self, data_loader, **kwargs):
-    # 遍历 dataset，共返回一个 epoch 的 batch 数据
-    for i, data_batch in enumerate(data_loader):
-        self.call_hook('before_train_iter')
-        # 验证时候 train_mode=False
-        self.run_iter(data_batch, train_mode=True, **kwargs)
-        self.call_hook('after_train_iter')
-   self.call_hook('after_train_epoch')
-```
-
-### IterBasedRunner
-
-不同于 `EpochBasedRunner`，`IterBasedRunner` 是指以 iter 为周期的工作流，例如设置 workflow = \[('train', 2)， ('val', 1)\] 表示循环迭代的训练 2 个 iter，然后验证 1 个 iter，MMSegmentation 语义分割框架默认采用的是  `IterBasedRunner`。
-
-其抽象逻辑如下所示：
-
-```python
-# 虽然是 iter 单位，但是某些场合需要 epoch 信息，由 IterLoader 提供
-iter_loaders = [IterLoader(x) for x in data_loaders]
-# 训练终止条件
-while curr_iter < max_iters:
-    # 遍历用户设置的工作流，例如 workflow = [('train', 2)， ('val', 1)]
-    for i, flow in enumerate(workflow):
-        # mode 是工作流函数，例如 train, iters 是迭代次数
-        mode, iters = flow
-        # 要么调用 self.train()，要么调用 self.val()
-        iter_runner = getattr(self, mode)
-        # 运行对应工作流函数
-        for _ in range(iters):
-            iter_runner(iter_loaders[i], **kwargs)
-```
-
-目前支持训练和验证两个工作流，以验证函数为例，其抽象逻辑是：
-
-```python
-# iter_runner 目前可以是 train 或者 val
-def val(self, data_loader, **kwargs):
-    # 获取 batch 数据，用于一次迭代
-    data_batch = next(data_loader)
-    self.call_hook('before_val_iter')
-    outputs = self.model.val_step(data_batch, self.optimizer, **kwargs)
-    self.outputs = outputs
-    self.call_hook('after_val_iter')
-```
-
-除了上述基础功能外，`EpochBasedRunner` 和 `IterBasedRunner` 还提供了 resume 、 save_checkpoint 和注册 hook 功能。
-
-### 一个简单例子
-
-以最常用的分类任务为例详细说明 `runner` 的使用方法。 开启任何一个训练任务，都需要包括如下步骤：
-
-**(1) dataloader、model 和优化器等类初始化**
-
-```python
-# 模型类初始化
-model=...
-# 优化器类初始化，典型值 cfg.optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0001)
-optimizer = build_optimizer(model, cfg.optimizer)
-# 工作流对应的 dataloader 初始化
-data_loaders = [
-        build_dataloader(
-            ds,
-            cfg.data.samples_per_gpu,
-            cfg.data.workers_per_gpu,
-            ...) for ds in dataset
-    ]
-```
-
-**(2) runner 类初始化**
-
-```python
-runner = build_runner(
-    # cfg.runner 典型配置为
-    # runner = dict(type='EpochBasedRunner', max_epochs=200)
-    cfg.runner,
-    default_args=dict(
-        model=model,
-        batch_processor=None,
-        optimizer=optimizer,
-        logger=logger))
-```
-
-**(3) 注册默认训练所必须的 hook，和用户自定义 hook**
-
-```python
-# 注册定制必需的 hook
-runner.register_training_hooks(
-    # lr相关配置，典型为
-    # lr_config = dict(policy='step', step=[100, 150])
-    cfg.lr_config,
-    # 优化相关配置，例如 grad_clip 等
-    optimizer_config,
-    # 权重保存相关配置，典型为
-    # checkpoint_config = dict(interval=1)，每个单位都保存权重
-    cfg.checkpoint_config,
-    # 日志相关配置
-    cfg.log_config,
-    ...)
-
-# 注册用户自定义 hook
-# 例如想使用 ema 功能，则可以设置 custom_hooks=[dict(type='EMAHook')]
-if cfg.get('custom_hooks', None):
-    custom_hooks = cfg.custom_hooks
-    for hook_cfg in cfg.custom_hooks:
-        hook_cfg = hook_cfg.copy()
-        priority = hook_cfg.pop('priority', 'NORMAL')
-        hook = build_from_cfg(hook_cfg, HOOKS)
-        runner.register_hook(hook, priority=priority)
-```
-
-然后可以进行 resume 或者 load_checkpoint 对权重进行加载。
-
-**(4) 开启训练流**
-
-```python
-# workflow 典型为 workflow = [('train', 1)]
-# 此时就真正开启了训练
-runner.run(data_loaders, cfg.workflow)
-```
-
-关于 workflow 设置，以 `EpochBasedRunner` 为例，详情如下：
-
-- 假设只想运行训练工作流，则可以设置 workflow = \[('train', 1)\]，表示只进行迭代训练
-- 假设想运行训练和验证工作流，则可以设置 workflow = \[('train',  3), ('val', 1)\]，表示先训练 3 个 epoch ，然后切换到 val 工作流，运行 1 个 epoch，然后循环，直到训练 epoch 次数达到指定值
-- 工作流设置还自由定制，例如你可以先验证再训练 workflow = \[('val', 1), ('train', 1)\]
-
-上述代码都已经封装到了各个代码库的 train.py 中，用户只需要设置相应的配置即可，上述流程会自动运行。

examples/train.py

-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.optim as optim
-import torchvision.transforms as transforms
-from torch.utils.data import DataLoader
-from torchvision.datasets import CIFAR10
-
-from mmcv.parallel import MMDataParallel
-from mmcv.runner import EpochBasedRunner
-from mmcv.utils import get_logger
-
-
-class Model(nn.Module):
-
-    def __init__(self):
-        super().__init__()
-        self.conv1 = nn.Conv2d(3, 6, 5)
-        self.pool = nn.MaxPool2d(2, 2)
-        self.conv2 = nn.Conv2d(6, 16, 5)
-        self.fc1 = nn.Linear(16 * 5 * 5, 120)
-        self.fc2 = nn.Linear(120, 84)
-        self.fc3 = nn.Linear(84, 10)
-        self.loss_fn = nn.CrossEntropyLoss()
-
-    def forward(self, x):
-        x = self.pool(F.relu(self.conv1(x)))
-        x = self.pool(F.relu(self.conv2(x)))
-        x = x.view(-1, 16 * 5 * 5)
-        x = F.relu(self.fc1(x))
-        x = F.relu(self.fc2(x))
-        x = self.fc3(x)
-        return x
-
-    def train_step(self, data, optimizer):
-        images, labels = data
-        predicts = self(images)  # -> self.__call__() -> self.forward()
-        loss = self.loss_fn(predicts, labels)
-        return {'loss': loss}
-
-
-if __name__ == '__main__':
-    model = Model()
-    if torch.cuda.is_available():
-        # only use gpu:0 to train
-        # Solved issue https://github.com/open-mmlab/mmcv/issues/1470
-        model = MMDataParallel(model.cuda(), device_ids=[0])
-
-    # dataset and dataloader
-    transform = transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
-    ])
-    trainset = CIFAR10(
-        root='data', train=True, download=True, transform=transform)
-    trainloader = DataLoader(
-        trainset, batch_size=128, shuffle=True, num_workers=2)
-
-    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
-    logger = get_logger('mmcv')
-    # runner is a scheduler to manage the training
-    runner = EpochBasedRunner(
-        model,
-        optimizer=optimizer,
-        work_dir='./work_dir',
-        logger=logger,
-        max_epochs=4)
-
-    # learning rate scheduler config
-    lr_config = dict(policy='step', step=[2, 3])
-    # configuration of optimizer
-    optimizer_config = dict(grad_clip=None)
-    # configuration of saving checkpoints periodically
-    checkpoint_config = dict(interval=1)
-    # save log periodically and multiple hooks can be used simultaneously
-    log_config = dict(interval=100, hooks=[dict(type='TextLoggerHook')])
-    # register hooks to runner and those hooks will be invoked automatically
-    runner.register_training_hooks(
-        lr_config=lr_config,
-        optimizer_config=optimizer_config,
-        checkpoint_config=checkpoint_config,
-        log_config=log_config)
-
-    runner.run([trainloader], [('train', 1)])

mmcv/__init__.py

@@ -10,7 +10,4 @@ from .visualization import *
 
 # The following modules are not imported to this level, so mmcv may be used
 # without PyTorch.
-# - runner
-# - parallel
 # - op
-# - device

mmcv/cnn/alexnet.py

@@ -4,6 +4,7 @@ from typing import Optional
 
 import torch
 import torch.nn as nn
+from mmengine.runner import load_checkpoint
 
 
 class AlexNet(nn.Module):
@@ -45,7 +46,6 @@ class AlexNet(nn.Module):
     def init_weights(self, pretrained: Optional[str] = None) -> None:
         if isinstance(pretrained, str):
             logger = logging.getLogger()
-            from ..runner import load_checkpoint
             load_checkpoint(self, pretrained, strict=False, logger=logger)
         elif pretrained is None:
             # use default initializer

mmcv/cnn/resnet.py

@@ -5,6 +5,7 @@ from typing import Optional, Sequence, Tuple, Union
 import torch.nn as nn
 import torch.utils.checkpoint as cp
 from mmengine.model.utils import constant_init, kaiming_init
+from mmengine.runner import load_checkpoint
 from torch import Tensor
 
 
@@ -270,7 +271,6 @@ class ResNet(nn.Module):
     def init_weights(self, pretrained: Optional[str] = None) -> None:
         if isinstance(pretrained, str):
             logger = logging.getLogger()
-            from ..runner import load_checkpoint
             load_checkpoint(self, pretrained, strict=False, logger=logger)
         elif pretrained is None:
             for m in self.modules():

mmcv/cnn/vgg.py

@@ -4,6 +4,7 @@ from typing import List, Optional, Sequence, Tuple, Union
 
 import torch.nn as nn
 from mmengine.model.utils import constant_init, kaiming_init, normal_init
+from mmengine.runner import load_checkpoint
 from torch import Tensor
 
 
@@ -126,7 +127,6 @@ class VGG(nn.Module):
     def init_weights(self, pretrained: Optional[str] = None) -> None:
         if isinstance(pretrained, str):
             logger = logging.getLogger()
-            from ..runner import load_checkpoint
             load_checkpoint(self, pretrained, strict=False, logger=logger)
         elif pretrained is None:
             for m in self.modules():

mmcv/device/__init__.py

-# Copyright (c) OpenMMLab. All rights reserved.
-from . import ipu, mlu, mps
-from .scatter_gather import scatter, scatter_kwargs
-from .utils import get_device
-
-__all__ = ['mlu', 'ipu', 'mps', 'get_device', 'scatter', 'scatter_kwargs']

mmcv/device/_functions.py

-# Copyright (c) OpenMMLab. All rights reserved.
-from typing import List, Union
-
-import torch
-
-from mmcv.utils import deprecated_api_warning
-from .utils import get_device
-
-
-def scatter(input: Union[List, torch.Tensor], devices: List) -> List:
-    """scatter copies tensor to devices directly."""
-    current_device = get_device()
-    if isinstance(input, list):
-        outputs = [scatter(_input, devices) for _input in input]
-        return outputs
-    elif isinstance(input, torch.Tensor):
-        output = input.contiguous()
-        return output.to(current_device) if devices != [-1] else output
-    else:
-        raise Exception(f'Unknown type {type(input)}.')
-
-
-class Scatter:
-
-    @staticmethod
-    @deprecated_api_warning({'target_mlus': 'target_devices'},
-                            cls_name='Scatter')
-    def forward(target_devices, input):
-        outputs = scatter(input, target_devices)
-        return tuple(outputs) if isinstance(outputs, list) else (outputs, )

mmcv/device/ipu/__init__.py

-# Copyright (c) OpenMMLab. All rights reserved.
-from mmcv.utils import IS_IPU_AVAILABLE
-
-if IS_IPU_AVAILABLE:
-    from .dataloader import IPUDataLoader
-    from .hook_wrapper import IPUFp16OptimizerHook
-    from .model_wrapper import ipu_model_wrapper
-    from .runner import IPUBaseRunner, IPUEpochBasedRunner, IPUIterBasedRunner
-    from .utils import cfg2options
-    __all__ = [
-        'cfg2options', 'ipu_model_wrapper', 'IPUFp16OptimizerHook',
-        'IPUDataLoader', 'IPUBaseRunner', 'IPUEpochBasedRunner',
-        'IPUIterBasedRunner'
-    ]

mmcv/device/ipu/dataloader.py

-# Copyright (c) OpenMMLab. All rights reserved.
-from collections.abc import Mapping, Sequence
-from functools import partial
-
-import poptorch
-from torch.utils.data.dataloader import default_collate
-
-from mmcv.parallel import DataContainer
-
-
-def collate(batch, samples_per_gpu=1):
-    """Put each data field into a tensor/DataContainer with outer dimension
-    batch size.
-
-    TODO support for
-    :type:`~mmcv.parallel.DataContainer`. Currently, it will be ignored.
-    There are 3 cases.
-
-    1. cpu_only = True, e.g., meta data.
-    2. cpu_only = False, stack = True, e.g., images tensors.
-    3. cpu_only = False, stack = False, e.g., gt bboxes.
-    """
-
-    if not isinstance(batch, Sequence):
-        raise TypeError(
-            f'`batch` should be a sequence, but got {type(batch)}.')
-
-    if isinstance(batch[0], DataContainer):
-        # TODO `DataContainer` will be supported in the future.
-        raise TypeError('DataContainer is not supported in ipu data loader.')
-    elif isinstance(batch[0], Sequence):
-        transposed = zip(*batch)
-        collated_batch = []
-        for samples in transposed:
-            if not isinstance(samples[0], DataContainer):
-                # At present, we will skip the processing of datacontainer,
-                # which will reduce the performance of IPU DataLoder
-                collated_batch.append(collate(samples, samples_per_gpu))
-        return collated_batch
-    elif isinstance(batch[0], Mapping):
-        collated_batch = {}
-        for key in batch[0]:
-            if not isinstance(batch[0][key], DataContainer):
-                # At present, we will skip the processing of datacontainer,
-                # which will reduce the performance of IPU DataLoder
-                collated_batch[key] = collate([d[key] for d in batch])
-        return collated_batch
-    else:
-        return default_collate(batch)
-
-
-class IPUDataLoader(poptorch.DataLoader):
-    """Thin wrapper of `torch.utils.data.DataLoader`.
-
-    Compared with the pytorch DataLoder, this DataLoder changes the way of
-    calculation of batch size and adds the AsynchronousDataAccessor to
-    load and release data faster in cpu mode.
-
-    If this data loader is used in a distributed execution environment, it will
-    ensure that each process uses a different subset of the dataset, providing
-    you first call ``options.randomSeed(N)`` with an integer N which is the
-    same across all hosts.
-
-    Args:
-        dataset (torch.utils.data.Dataset): The dataset to get the data from.
-        options (poptorch.Options): Options that will be used to compile
-            and run the model.
-        batch_size (int, optional): This is the batch size in the conventional
-            sense of being the size that runs through an operation in the model
-            at any given time.
-        shuffle (bool, optional): set to ``True`` to have the data reshuffled
-            at every epoch (default: ``False``).
-        num_workers (int, optional): how many subprocesses to use for data
-            loading. ``0`` means that the data will be loaded in the main
-            process. (default: ``0``)
-        drop_last (bool, optional): If True and the number of elements in the
-            dataset is not a multiple of the combined batch size then the
-            incomplete batch at the end will be dropped.
-        persistent_workers (bool, optional): Re-use workers between
-            iterations if True.
-        auto_distributed_partitioning (bool, optional): If True, partitions the
-            dataset for distributed execution automatically. Otherwise, it is
-            assumed that partitioning has been handled manually.
-        mode (poptorch.DataLoaderMode, optional): If `DataLoaderMode.Async`,
-            uses an :py:class:`~poptorch.AsynchronousDataAccessor` to access
-            the dataset. If `DataLoaderMode.Sync`, accesses the dataset
-            synchronously.
-        async_options (Dict[str, Any], optional): Options to pass to
-            :py:class:`~poptorch.AsynchronousDataAccessor`.
-        rebatched_worker_size (int, optional): When using AsyncRebatched: batch
-            size of the tensors loaded by the workers.
-            Default to the combined batch size.
-            If specified the ``rebatched_worker_size`` must be less than
-            or equal to the combined batch size.
-        kwargs (Dict[str, Any], optional): Other options to pass to PyTorch's
-            ``DataLoader`` constructor.
-    """
-
-    def __init__(self,
-                 dataset,
-                 options,
-                 batch_size=1,
-                 shuffle=False,
-                 num_workers=0,
-                 drop_last=True,
-                 persistent_workers=True,
-                 auto_distributed_partitioning=True,
-                 mode='sync',
-                 async_options=None,
-                 rebatched_worker_size=None,
-                 **kwargs):
-        """Lazy init:
-
-        In many frameworks, the dataloader will be constructed before the
-        initialization of the ipu options, so the lazy init method is used
-        here, and the real initialization will not be done until the dataloader
-        needs to be used and the options are input.
-        """
-        # lazy init: sometimes, we cannot get IPU options when build data
-        #            loader
-        self.kwargs = {
-            'dataset': dataset,
-            'batch_size': batch_size,
-            'shuffle': shuffle,
-            'num_workers': num_workers,
-            'drop_last': drop_last,
-            'persistent_workers': persistent_workers,
-            'auto_distributed_partitioning': auto_distributed_partitioning,
-            'mode': mode,
-            'collate_fn': partial(collate, samples_per_gpu=batch_size),
-            'async_options': async_options,
-            'rebatched_worker_size': rebatched_worker_size,
-            **kwargs
-        }
-        self.dataset = dataset
-        self.initialized = False
-        if options:
-            self.init(options=options)
-
-    def init(self, options, **kwargs):
-        if not self.initialized:
-            kwargs = {**self.kwargs, **kwargs, 'options': options}
-            if kwargs['mode'] == 'sync':
-                kwargs['mode'] = poptorch.DataLoaderMode.Sync
-            elif kwargs['mode'] == 'async':
-                kwargs['mode'] = poptorch.DataLoaderMode.AsyncRebatched
-                if kwargs['async_options'] is None:
-                    kwargs['async_options'] = {
-                        'load_indefinitely': True,
-                        'buffer_size': 8
-                    }
-                if kwargs['rebatched_worker_size'] is None:
-                    kwargs['rebatched_worker_size'] = 128
-            super().__init__(**kwargs)
-            self.initialized = True
-
-        return self

mmcv/device/ipu/hierarchical_data_manager.py

-# Copyright (c) OpenMMLab. All rights reserved.
-import warnings
-
-import numpy as np
-import torch
-
-from mmcv.parallel import DataContainer
-
-# A customized None type for HierarchicalDataManager
-HierarchicalDataNone = object()
-
-
-class HierarchicalDataManager:
-    """A class manage all the tensors in the hierarchical data.
-
-    At present, the input data structure accepted by IPU is limited,
-    when the input data structure of mmcv varies.
-    Here, an intermediate class is needed to get and update tensors
-    from the original data.
-
-    HierarchicalDataManager will record a hierarchical input/output data in
-    self._hierarchical_data. For example, we have an input data:
-    {'img': tensorA, 'label': tensorB, 'img_metas': [tensorC, tensorD]}
-    To enable IPU to use the input, HierarchicalDataManager will collect
-    the torch tensors from self._hierarchical_data into a tuple like:
-    (tensorA, tensorB, tensorC, tensorD).
-    Meanwhile, the return of IPU is a tuple of tensors, HierarchicalDataManager
-    also have a function named update_all_tensors to update tensors in
-    self._hierarchical_data which is the output for upper calls.
-
-    Args:
-        logger (:obj:`logging.Logger`): Logger used during running.
-             Defaults to None.
-    """
-
-    def __init__(self, logger=None):
-        self.atomic_types = (int, str, float, np.ndarray, type(None))
-        self.warning = warnings.warn if logger is None else logger.warning
-        # enable or disable input data's shape and value check
-        self.quick_mode = False
-        self._hierarchical_data = None
-
-    def quick(self):
-        self.quick_mode = True
-
-    def compare_atomic_type(self, a, b):
-        """Compare data, supported datatypes are numpy array and python basic
-        types."""
-        if isinstance(a, np.ndarray):
-            return np.all(a == b)
-        else:
-            return a == b
-
-    def record_hierarchical_data(self, data):
-        """Record a hierarchical data."""
-        if self._hierarchical_data is not None:
-            if isinstance(data, torch.Tensor):
-                assert isinstance(self._hierarchical_data, torch.Tensor), \
-                    'original hierarchical data is not torch.tensor'
-                self._hierarchical_data = data
-            else:
-                self.update_hierarchical_data(data)
-        else:
-            self._hierarchical_data = data
-
-    @property
-    def hierarchical_data(self):
-        return self._hierarchical_data
-
-    def update_hierarchical_data(self,
-                                 dataA,
-                                 dataB=HierarchicalDataNone,
-                                 strict=True,
-                                 address='data'):
-        """Update dataB with dataA in-place.
-
-        Args:
-            dataA (list or dict or tuple): New hierarchical data.
-            dataB (list or dict or tuple): hierarchical data to update.
-                if not specified, self.hierarchical_data will be updated then.
-            strict (bool, optional): If true, an error will be reported
-                when the following conditions occur:
-                1. Non-torch.Tensor data changed.
-                2. Torch.Tensor data shape changed.
-            address (str): Record the address of current data to be updated.
-                Default: 'data'.
-        """
-        if dataB is HierarchicalDataNone:
-            dataB = self.hierarchical_data
-
-        # Update with a da ta with the same structure
-        # but different values(tensors and basic python data types)
-        if isinstance(dataA, (tuple, list)):
-            for idx, node in enumerate(dataA):
-                new_address = ''
-                if not self.quick_mode:
-                    new_address = address + f'[{str(idx)}]'
-                    assert isinstance(node, type(dataB[idx])),\
-                        f'data structure changed: {new_address}'
-                if isinstance(node, torch.Tensor):
-                    dataB[idx] = node
-                else:
-                    self.update_hierarchical_data(
-                        node, dataB[idx], strict, address=new_address)
-        elif isinstance(dataA, dict):
-            for k, v in dataA.items():
-                new_address = ''
-                if not self.quick_mode:
-                    new_address = address + f'[{str(k)}]'
-                    assert isinstance(v, type(dataB[k])),\
-                        f'data structure changed: {new_address}'
-                if isinstance(v, torch.Tensor):
-                    dataB[k] = v
-                else:
-                    self.update_hierarchical_data(
-                        v, dataB[k], strict, address=new_address)
-        elif isinstance(dataA, self.atomic_types):
-            if not self.quick_mode:
-                is_equal = self.compare_atomic_type(dataA, dataB)
-                if not is_equal:
-                    if strict:
-                        raise ValueError(
-                            'all data except torch.Tensor should be same, '
-                            f'but data({address}) is changed.')
-                    else:
-                        self.warning(
-                            f'find a non-torch.Tensor data({type(dataA)}) '
-                            f'changed, and the address is {address}')
-        elif isinstance(dataA, DataContainer):
-            if not self.quick_mode:
-                assert isinstance(dataB, DataContainer)
-                new_address = address + '.data'
-                self.update_hierarchical_data(
-                    dataA.data, dataB.data, False, address=new_address)
-        else:
-            raise NotImplementedError(
-                f'not supported datatype:{type(dataA)}, address is {address}')
-
-    def collect_all_tensors(self, hierarchical_data=None):
-        """Collect torch.Tensor data from self.hierarchical_data to a list and
-        return."""
-        # get a list of tensor from self._hierarchical_data
-        if hierarchical_data is None:
-            hierarchical_data = self._hierarchical_data
-        tensors = []
-        if isinstance(hierarchical_data, torch.Tensor):
-            tensors = [hierarchical_data]
-        else:
-            self._collect_tensors(hierarchical_data, tensors)
-        return tensors
-
-    def _collect_tensors(self, data, tensors):
-        if isinstance(data, (tuple, list)):
-            for node in data:
-                if isinstance(node, torch.Tensor):
-                    tensors.append(node)
-                else:
-                    self._collect_tensors(node, tensors)
-        elif isinstance(data, dict):
-            for v in data.values():
-                if isinstance(v, torch.Tensor):
-                    tensors.append(v)
-                else:
-                    self._collect_tensors(v, tensors)
-        elif isinstance(data, self.atomic_types):
-            pass
-        elif isinstance(data, DataContainer):
-            self._collect_tensors(data.data, tensors)
-        else:
-            raise NotImplementedError(f'not supported datatype:{type(data)}')
-
-    def update_all_tensors(self, tensors):
-        """Put tensors from tuple back to self.hierarchical_data."""
-        if isinstance(self._hierarchical_data, torch.Tensor):
-            print(tensors, len(tensors))
-            assert len(tensors) == 1
-            assert isinstance(tensors[0], torch.Tensor)
-            self._hierarchical_data = tensors[0]
-        else:
-            # convert to list if tensors is tuple
-            tensors = list(tensors)
-            self._set_tensors(self._hierarchical_data, tensors)
-        return self.hierarchical_data
-
-    def _set_tensors(self, data, tensors):
-        if isinstance(data, tuple):
-            data = list(data)
-            for idx in range(len(data)):
-                if isinstance(data[idx], torch.Tensor):
-                    data[idx] = tensors.pop(0)
-                else:
-                    self._set_tensors(data[idx], tensors)
-            data = tuple(data)
-        elif isinstance(data, list):
-            for idx in range(len(data)):
-                if isinstance(data[idx], torch.Tensor):
-                    data[idx] = tensors.pop(0)
-                else:
-                    self._set_tensors(data[idx], tensors)
-        elif isinstance(data, dict):
-            for k, v in data.items():
-                if isinstance(v, torch.Tensor):
-                    data[k] = tensors.pop(0)
-                else:
-                    self._set_tensors(v, tensors)
-        elif isinstance(data, self.atomic_types):
-            pass
-        elif isinstance(data, DataContainer):
-            self._set_tensors(data.data, tensors)
-        else:
-            raise NotImplementedError(f'not supported datatype:{type(data)}')
-
-    def clean_all_tensors(self):
-        """Delete tensors from self.hierarchical_data."""
-        self._clean_tensors(self._hierarchical_data)
-
-    def _clean_tensors(self, data):
-        if isinstance(data, tuple):
-            data = list(data)
-            for idx in range(len(data)):
-                if isinstance(data[idx], torch.Tensor):
-                    data[idx] = None
-                else:
-                    self._clean_tensors(data[idx])
-            data = tuple(data)
-        elif isinstance(data, list):
-            for idx in range(len(data)):
-                if isinstance(data[idx], torch.Tensor):
-                    data[idx] = None
-                else:
-                    self._clean_tensors(data[idx])
-        elif isinstance(data, dict):
-            for k, v in data.items():
-                if isinstance(v, torch.Tensor):
-                    data[k] = None
-                else:
-                    self._clean_tensors(v)
-        elif isinstance(data, self.atomic_types):
-            pass
-        elif isinstance(data, DataContainer):
-            self._clean_tensors(data.data)
-        else:
-            raise NotImplementedError(f'not supported datatype:{type(data)}')

mmcv/device/ipu/hook_wrapper.py

-# Copyright (c) OpenMMLab. All rights reserved.
-from mmcv.runner import HOOKS, LrUpdaterHook, OptimizerHook
-from mmcv.utils import TORCH_VERSION, digit_version
-
-
-def wrap_lr_updater_hook(lr_hook_class):
-    """A wrapper function to wrap any subclass of LrUpdaterHook.
-
-    IPU needs extra operations to upload optimizer settings. This wrapper will
-    override function(_set_lr) of a subclass of LrUpdaterHook.
-    """
-    assert issubclass(lr_hook_class, LrUpdaterHook)
-
-    class ipu_lr_hook_class(lr_hook_class):
-
-        def _set_lr(self, runner, *args, **kwargs):
-            super()._set_lr(runner, *args, **kwargs)
-            # convert torch optimizer to poptorch optimizer
-            runner.model.setOptimizer(runner.optimizer)
-
-    return ipu_lr_hook_class
-
-
-def wrap_optimizer_hook(optimizer_hook_class):
-    """A wrapper function to wrap OptimizerHook.
-
-    This is an non-intrusive implementation of wrapping optimizer hook (or you
-    need to change every config file to use IPU optimizer hook) IPU's clip-norm
-    implementation is different from pytorch, so there should be an error
-    raised when using clip-norm.
-    """
-
-    class ipu_optimizer_hook_class(OptimizerHook):
-
-        def __init__(self, **kwargs):
-            super().__init__(**kwargs)
-            if self.grad_clip is not None:
-                raise NotImplementedError('IPU does not support gradient clip')
-
-    return ipu_optimizer_hook_class
-
-
-if (TORCH_VERSION != 'parrots'
-        and digit_version(TORCH_VERSION) >= digit_version('1.6.0')):
-
-    @HOOKS.register_module()
-    class IPUFp16OptimizerHook(OptimizerHook):
-        """FP16 optimizer hook (using PyTorch's implementation).
-
-        If you are using PyTorch >= 1.6, torch.cuda.amp is used as the backend,
-        to take care of the optimization procedure.
-
-        Args:
-            loss_scale (float | str | dict): Scale factor configuration.
-                If loss_scale is a float, static loss scaling will be used with
-                the specified scale. If loss_scale is a string, it must be
-                'dynamic', then dynamic loss scaling will be used.
-                It can also be a dict containing arguments of GradScalar.
-                Defaults to 512. For Pytorch >= 1.6, mmcv uses official
-                implementation of GradScaler. If you use a dict version of
-                loss_scale to create GradScaler, please refer to:
-                https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler
-                for the parameters.
-
-        Examples:
-            >>> loss_scale = dict(
-            ...     init_scale=65536.0,
-            ...     growth_factor=2.0,
-            ...     backoff_factor=0.5,
-            ...     growth_interval=2000
-            ... )
-            >>> optimizer_hook = Fp16OptimizerHook(loss_scale=loss_scale)
-        """
-
-        def __init__(self,
-                     grad_clip=None,
-                     coalesce=True,
-                     bucket_size_mb=-1,
-                     loss_scale=512.,
-                     distributed=True):
-            assert grad_clip is None,\
-                'IPU mode does not support `grad_clip` currently'
-            assert coalesce,\
-                'implemented all reduce in distributed training currently'
-            assert bucket_size_mb == -1,\
-                '`bucket_size_mb` should not be set in IPU mode'
-            self.distributed = distributed
-            self._scale_update_param = None
-            if loss_scale == 'dynamic':
-                raise NotImplementedError(
-                    'IPU mode does not support dynamic loss scale currently')
-            elif isinstance(loss_scale, float):
-                self.loss_scale = loss_scale
-            elif isinstance(loss_scale, dict):
-                raise NotImplementedError(
-                    'IPU mode supports single scale currently')
-            else:
-                raise ValueError(
-                    f'loss_scale should be float, but got {loss_scale} ')
-
-        def after_train_iter(self, runner):
-            pass
-
-else:
-    raise RuntimeError('The IPU mode only supports torch 1.6 and above')

mmcv/device/ipu/runner.py

-# Copyright (c) OpenMMLab. All rights reserved.
-
-from mmcv.runner import (HOOKS, RUNNERS, BaseRunner, EpochBasedRunner,
-                         IterBasedRunner)
-from mmcv.utils import IS_IPU_AVAILABLE
-
-if IS_IPU_AVAILABLE:
-    from .dataloader import IPUDataLoader
-    from .hook_wrapper import (IPUFp16OptimizerHook, wrap_lr_updater_hook,
-                               wrap_optimizer_hook)
-    from .model_wrapper import ipu_model_wrapper
-    from .utils import build_from_cfg_with_wrapper, cfg2options
-
-
-class IPUBaseRunner(BaseRunner):
-    """A base runner for IPU.
-
-    This runner has some extra processes for IPU which are shown below:
-
-    1. Parse options for IPU
-    2. wrap pytorch model for IPU
-    3. Raise errors while encountering illegal usage
-    4. Input IPU options and initialize dataloader if finding an instance
-       of IPUDataLoader
-
-    Args:
-        model (:obj:`nn.Module`): The model to run.
-        options_cfg (mmcv.Config, dict): Options that will be used to compile
-            and run the model.
-        modules_to_record (mmcv.Config, list): Index or name of modules which
-            will be recorded for output. It is necessary to specify output for
-            static graph of model training or inference.
-        ipu_model_cfg (mmcv.Config, dict): Config of model partition and
-            recomputing checkpoint
-        fp16_cfg (mmcv.Config): Config for fp16 training.
-        batch_processor (callable): A callable method that process a data
-            batch. Should be None for IPU runner
-        kwargs (Dict[str, Any], optional): Keyword arguments will be passed to
-        ``base_runner.BaseRunner``.
-    """
-
-    def __init__(self,
-                 model,
-                 options_cfg=None,
-                 modules_to_record=None,
-                 ipu_model_cfg=None,
-                 fp16_cfg=None,
-                 batch_processor=None,
-                 **kwargs):
-        assert hasattr(model, 'train_step') and batch_processor is None,\
-            'only support model with train_step'
-
-        if options_cfg is None:
-            options_cfg = {}
-        # call BaseRunner.__init__() here
-        super().__init__(model, **kwargs)
-
-        # process options of ipu
-        if IS_IPU_AVAILABLE:
-            self.options = cfg2options(options_cfg)
-            self.model = ipu_model_wrapper(
-                self.model,
-                self.options,
-                self.optimizer,
-                self.logger,
-                modules_to_record=modules_to_record,
-                ipu_model_cfg=ipu_model_cfg,
-                fp16_cfg=fp16_cfg)
-        else:
-            raise NotImplementedError('cpu mode on IPURunner is not supported')
-
-    def register_lr_hook(self, lr_config):
-        if lr_config is None:
-            return
-        assert isinstance(lr_config, dict)
-        assert 'policy' in lr_config
-        policy_type = lr_config.pop('policy')
-        # If the type of policy is all in lower case,
-        # e.g., 'cyclic', then its first letter will be capitalized,
-        # e.g., to be 'Cyclic'.
-        # This is for the convenient usage of Lr updater.
-        # Since this is not applicable for `
-        # CosineAnnealingLrUpdater`, the string will not be changed
-        # if it contains capital letters.
-        if policy_type == policy_type.lower():
-            policy_type = policy_type.title()
-        hook_type = policy_type + 'LrUpdaterHook'
-        lr_config['type'] = hook_type
-        hook = build_from_cfg_with_wrapper(lr_config, HOOKS,
-                                           wrap_lr_updater_hook)
-        self.register_hook(hook, priority='VERY_HIGH')
-
-    def register_optimizer_hook(self, optimizer_config):
-        if optimizer_config is None:
-            return
-        assert isinstance(optimizer_config, (dict, IPUFp16OptimizerHook))
-        if isinstance(optimizer_config, dict):
-            optimizer_config.setdefault('type', 'OptimizerHook')
-            hook = build_from_cfg_with_wrapper(optimizer_config, HOOKS,
-                                               wrap_optimizer_hook)
-        else:
-            hook = optimizer_config
-        self.register_hook(hook, priority='ABOVE_NORMAL')
-
-    def run(self, data_loaders, workflow, *args, **kwargs):
-        for i, flow in enumerate(workflow):
-            mode, _ = flow
-            # initialize IPU dataloader if not initialized
-            assert isinstance(data_loaders[i], IPUDataLoader),\
-                'IPU runner can only work with `IPUDataLoader`'
-            data_loaders[i].init(options=self.get_options(mode))
-
-        super().run(data_loaders, workflow, *args, **kwargs)
-
-    def get_options(self, mode):
-        if mode == 'train':
-            return self.options['training']
-        elif mode == 'val':
-            return self.options['inference']
-        else:
-            raise ValueError(f'mode should be train or val but got {mode}')
-
-
-@RUNNERS.register_module()
-class IPUEpochBasedRunner(IPUBaseRunner, EpochBasedRunner):
-    """Epoch-based Runner for IPU.
-
-    The Inheritance order(MRO) is: IPUEpochBasedRunner -> IPUBaseRunner ->
-    EpochBasedRunner -> BaseRunner This runner train models epoch by epoch.
-    """
-    pass
-
-
-@RUNNERS.register_module()
-class IPUIterBasedRunner(IPUBaseRunner, IterBasedRunner):
-    """Iteration-based Runner for IPU.
-
-    The Inheritance order(MRO) is: IPUIterBasedRunner -> IPUBaseRunner ->
-    IterBasedRunner -> BaseRunner This runner train models iteration by
-    iteration.
-    """
-    pass