Support CPU Train/Inference (#86)

docs/en/tutorials/1_inference.md

@@ -41,6 +41,7 @@ We provide testing scripts for evaluating an existing model on the whole dataset
 The following testing environments are supported:
 
 - single GPU
+- CPU
 - single node multiple GPUs
 - multiple nodes
 
@@ -55,6 +56,15 @@ python tools/test.py \
  [--out-dir ${OUTPUT_DIRECTORY}] \
  [--show-dir ${VISUALIZATION_DIRECTORY}]
 
+# CPU: disable GPUs and run single-gpu testing script
+export CUDA_VISIBLE_DEVICES=-1
+python tools/test.py \
+ ${CONFIG_FILE} \
+ ${CHECKPOINT_FILE} \
+ [--out ${RESULT_FILE}] \
+ [--eval ${EVAL_METRICS}] \
+ [--show]
+
 # multi-gpu testing
 bash tools/dist_test.sh \
  ${CONFIG_FILE} \

docs/en/tutorials/2_finetune.md

@@ -65,6 +65,20 @@ Difference between `resume-from` and `load-from`:
 It is usually used for resuming the training process that is interrupted accidentally.
 `load-from` only loads the model weights and the training iteration starts from 0. It is usually used for finetuning.
 
+### Training on CPU
+
+The process of training on the CPU is consistent with single GPU training. We just need to disable GPUs before the training process.
+
+```shell
+export CUDA_VISIBLE_DEVICES=-1
+```
+
+And then run the script [above](#training-on-a-single-GPU).
+
+```{note}
+We do not recommend users to use CPU for training because it is too slow. We support this feature to allow users to debug on machines without GPU for convenience.
+```
+
 ## Training on multiple GPUs
 
 MMFlow implements **distributed** training with `MMDistributedDataParallel`.

mmflow/apis/train.py

@@ -3,14 +3,17 @@
 import warnings
 from typing import Optional, Sequence, Union
 
+import mmcv
 import numpy as np
 import torch
 import torch.distributed as dist
+from mmcv.cnn.utils import revert_sync_batchnorm
 from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
 from mmcv.runner import (HOOKS, Fp16OptimizerHook, OptimizerHook,
  build_optimizer, build_runner, get_dist_info)
 from mmcv.utils import Config, build_from_cfg
 
+from mmflow import digit_version
 from mmflow.core import DistEvalHook, EvalHook
 from mmflow.datasets import build_dataloader, build_dataset
 from mmflow.utils import find_latest_checkpoint, get_root_logger
@@ -114,8 +117,17 @@ def train_model(model: Module,
  broadcast_buffers=False,
  find_unused_parameters=find_unused_parameters)
  else:
- model = MMDataParallel(
- model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)
+ # SyncBN is not support for DP
+ warnings.warn(
+ 'SyncBN is only supported with DDP. To be compatible with DP, '
+ 'we convert SyncBN to BN. Please use dist_train.sh which can '
+ 'avoid this error.')
+ model = revert_sync_batchnorm(model)
+ if not torch.cuda.is_available():
+ assert digit_version(mmcv.__version__) >= digit_version('1.4.4'), \
+ 'Please use MMCV >= 1.4.4 for CPU training!'
+
+ model = MMDataParallel(model, device_ids=cfg.gpu_ids)
 
  # build runner
  optimizer = build_optimizer(model, cfg.optimizer)

tools/test.py

@@ -1,16 +1,19 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import argparse
 import os
+import warnings
 
 import mmcv
 import torch
 from mmcv import Config, DictAction
 from mmcv.cnn import fuse_conv_bn
+from mmcv.cnn.utils import revert_sync_batchnorm
 from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
 from mmcv.runner import (get_dist_info, init_dist, load_checkpoint,
  wrap_fp16_model)
 from mmcv.utils.logging import print_log
 
+from mmflow import digit_version
 from mmflow.apis import multi_gpu_test, single_gpu_test
 from mmflow.core import online_evaluation
 from mmflow.datasets import build_dataloader, build_dataset
@@ -143,8 +146,18 @@ def main():
  load_checkpoint(model, args.checkpoint, map_location='cpu')
  if args.fuse_conv_bn:
  model = fuse_conv_bn(model)
+
  if not distributed:
- model = MMDataParallel(model, device_ids=[0])
+ warnings.warn(
+ 'SyncBN is only supported with DDP. To be compatible with DP, '
+ 'we convert SyncBN to BN. Please use dist_train.sh which can '
+ 'avoid this error.')
+ model = revert_sync_batchnorm(model)
+ if not torch.cuda.is_available():
+ assert digit_version(mmcv.__version__) >= digit_version('1.4.4'), \
+ 'Please use MMCV >= 1.4.4 for CPU training!'
+ model = MMDataParallel(model, device_ids=cfg.gpu_ids)
+
  else:
  model = MMDistributedDataParallel(
  model.cuda(),

tools/train.py

@@ -164,6 +164,7 @@ def main():
 
  model = build_flow_estimator(cfg.model)
  model.init_weights()
+
  logger.info(model)
  if cfg.data.train_dataloader.get('sample_ratio') is None:
  # build_dataset will concat the list of dataset