Merge branch 'add_Recommendation' into 'main'

添加openmmlab测试用例

See merge request dcutoolkit/deeplearing/dlexamples_new!32

openmmlab_test/mmclassification-speed-benchmark/mmcls/core/evaluation/eval_hooks.py

+import os.path as osp
+import warnings
+
+from mmcv.runner import Hook
+from torch.utils.data import DataLoader
+
+
+class EvalHook(Hook):
+    """Evaluation hook.
+
+    Args:
+        dataloader (DataLoader): A PyTorch dataloader.
+        interval (int): Evaluation interval (by epochs). Default: 1.
+    """
+
+    def __init__(self, dataloader, interval=1, by_epoch=True, **eval_kwargs):
+        warnings.warn(
+            'DeprecationWarning: EvalHook and DistEvalHook in mmcls will be '
+            'deprecated, please install mmcv through master branch.')
+        if not isinstance(dataloader, DataLoader):
+            raise TypeError('dataloader must be a pytorch DataLoader, but got'
+                            f' {type(dataloader)}')
+        self.dataloader = dataloader
+        self.interval = interval
+        self.eval_kwargs = eval_kwargs
+        self.by_epoch = by_epoch
+
+    def after_train_epoch(self, runner):
+        if not self.by_epoch or not self.every_n_epochs(runner, self.interval):
+            return
+        from mmcls.apis import single_gpu_test
+        results = single_gpu_test(runner.model, self.dataloader, show=False)
+        self.evaluate(runner, results)
+
+    def after_train_iter(self, runner):
+        if self.by_epoch or not self.every_n_iters(runner, self.interval):
+            return
+        from mmcls.apis import single_gpu_test
+        runner.log_buffer.clear()
+        results = single_gpu_test(runner.model, self.dataloader, show=False)
+        self.evaluate(runner, results)
+
+    def evaluate(self, runner, results):
+        eval_res = self.dataloader.dataset.evaluate(
+            results, logger=runner.logger, **self.eval_kwargs)
+        for name, val in eval_res.items():
+            runner.log_buffer.output[name] = val
+        runner.log_buffer.ready = True
+
+
+class DistEvalHook(EvalHook):
+    """Distributed evaluation hook.
+
+    Args:
+        dataloader (DataLoader): A PyTorch dataloader.
+        interval (int): Evaluation interval (by epochs). Default: 1.
+        tmpdir (str, optional): Temporary directory to save the results of all
+            processes. Default: None.
+        gpu_collect (bool): Whether to use gpu or cpu to collect results.
+            Default: False.
+    """
+
+    def __init__(self,
+                 dataloader,
+                 interval=1,
+                 gpu_collect=False,
+                 by_epoch=True,
+                 **eval_kwargs):
+        warnings.warn(
+            'DeprecationWarning: EvalHook and DistEvalHook in mmcls will be '
+            'deprecated, please install mmcv through master branch.')
+        if not isinstance(dataloader, DataLoader):
+            raise TypeError('dataloader must be a pytorch DataLoader, but got '
+                            f'{type(dataloader)}')
+        self.dataloader = dataloader
+        self.interval = interval
+        self.gpu_collect = gpu_collect
+        self.by_epoch = by_epoch
+        self.eval_kwargs = eval_kwargs
+
+    def after_train_epoch(self, runner):
+        if not self.by_epoch or not self.every_n_epochs(runner, self.interval):
+            return
+        from mmcls.apis import multi_gpu_test
+        results = multi_gpu_test(
+            runner.model,
+            self.dataloader,
+            tmpdir=osp.join(runner.work_dir, '.eval_hook'),
+            gpu_collect=self.gpu_collect)
+        if runner.rank == 0:
+            print('\n')
+            self.evaluate(runner, results)
+
+    def after_train_iter(self, runner):
+        if self.by_epoch or not self.every_n_iters(runner, self.interval):
+            return
+        from mmcls.apis import multi_gpu_test
+        runner.log_buffer.clear()
+        results = multi_gpu_test(
+            runner.model,
+            self.dataloader,
+            tmpdir=osp.join(runner.work_dir, '.eval_hook'),
+            gpu_collect=self.gpu_collect)
+        if runner.rank == 0:
+            print('\n')
+            self.evaluate(runner, results)

openmmlab_test/mmclassification-speed-benchmark/mmcls/core/evaluation/eval_metrics.py

+import numpy as np
+import torch
+
+
+def calculate_confusion_matrix(pred, target):
+    """Calculate confusion matrix according to the prediction and target.
+
+    Args:
+        pred (torch.Tensor | np.array): The model prediction with shape (N, C).
+        target (torch.Tensor | np.array): The target of each prediction with
+            shape (N, 1) or (N,).
+
+    Returns:
+        torch.Tensor: Confusion matrix with shape (C, C), where C is the number
+             of classes.
+    """
+
+    if isinstance(pred, np.ndarray):
+        pred = torch.from_numpy(pred)
+    if isinstance(target, np.ndarray):
+        target = torch.from_numpy(target)
+    assert (
+        isinstance(pred, torch.Tensor) and isinstance(target, torch.Tensor)), \
+        (f'pred and target should be torch.Tensor or np.ndarray, '
+         f'but got {type(pred)} and {type(target)}.')
+
+    num_classes = pred.size(1)
+    _, pred_label = pred.topk(1, dim=1)
+    pred_label = pred_label.view(-1)
+    target_label = target.view(-1)
+    assert len(pred_label) == len(target_label)
+    confusion_matrix = torch.zeros(num_classes, num_classes)
+    with torch.no_grad():
+        for t, p in zip(target_label, pred_label):
+            confusion_matrix[t.long(), p.long()] += 1
+    return confusion_matrix
+
+
+def precision_recall_f1(pred, target, average_mode='macro', thrs=None):
+    """Calculate precision, recall and f1 score according to the prediction and
+    target.
+
+    Args:
+        pred (torch.Tensor | np.array): The model prediction with shape (N, C).
+        target (torch.Tensor | np.array): The target of each prediction with
+            shape (N, 1) or (N,).
+        average_mode (str): The type of averaging performed on the result.
+            Options are 'macro' and 'none'. If 'none', the scores for each
+            class are returned. If 'macro', calculate metrics for each class,
+            and find their unweighted mean.
+            Defaults to 'macro'.
+        thrs (float | tuple[float], optional): Predictions with scores under
+            the thresholds are considered negative. Default to None.
+
+    Returns:
+        float | np.array | list[float | np.array]: Precision, recall, f1 score.
+            If the ``average_mode`` is set to macro, np.array is used in favor
+            of float to give class-wise results. If the ``average_mode`` is set
+             to none, float is used to return a single value.
+            If ``thrs`` is a single float or None, the function will return
+            float or np.array. If ``thrs`` is a tuple, the function will return
+             a list containing metrics for each ``thrs`` condition.
+    """
+
+    allowed_average_mode = ['macro', 'none']
+    if average_mode not in allowed_average_mode:
+        raise ValueError(f'Unsupport type of averaging {average_mode}.')
+
+    if isinstance(pred, torch.Tensor):
+        pred = pred.numpy()
+    if isinstance(target, torch.Tensor):
+        target = target.numpy()
+    assert (isinstance(pred, np.ndarray) and isinstance(target, np.ndarray)),\
+        (f'pred and target should be torch.Tensor or np.ndarray, '
+         f'but got {type(pred)} and {type(target)}.')
+
+    if thrs is None:
+        thrs = 0.0
+    if isinstance(thrs, float):
+        thrs = (thrs, )
+        return_single = True
+    elif isinstance(thrs, tuple):
+        return_single = False
+    else:
+        raise TypeError(
+            f'thrs should be float or tuple, but got {type(thrs)}.')
+
+    label = np.indices(pred.shape)[1]
+    pred_label = np.argsort(pred, axis=1)[:, -1]
+    pred_score = np.sort(pred, axis=1)[:, -1]
+
+    precisions = []
+    recalls = []
+    f1_scores = []
+    for thr in thrs:
+        # Only prediction values larger than thr are counted as positive
+        _pred_label = pred_label.copy()
+        if thr is not None:
+            _pred_label[pred_score <= thr] = -1
+        pred_positive = label == _pred_label.reshape(-1, 1)
+        gt_positive = label == target.reshape(-1, 1)
+        precision = (pred_positive & gt_positive).sum(0) / np.maximum(
+            pred_positive.sum(0), 1) * 100
+        recall = (pred_positive & gt_positive).sum(0) / np.maximum(
+            gt_positive.sum(0), 1) * 100
+        f1_score = 2 * precision * recall / np.maximum(precision + recall,
+                                                       1e-20)
+        if average_mode == 'macro':
+            precision = float(precision.mean())
+            recall = float(recall.mean())
+            f1_score = float(f1_score.mean())
+        precisions.append(precision)
+        recalls.append(recall)
+        f1_scores.append(f1_score)
+
+    if return_single:
+        return precisions[0], recalls[0], f1_scores[0]
+    else:
+        return precisions, recalls, f1_scores
+
+
+def precision(pred, target, average_mode='macro', thrs=None):
+    """Calculate precision according to the prediction and target.
+
+    Args:
+        pred (torch.Tensor | np.array): The model prediction with shape (N, C).
+        target (torch.Tensor | np.array): The target of each prediction with
+            shape (N, 1) or (N,).
+        average_mode (str): The type of averaging performed on the result.
+            Options are 'macro' and 'none'. If 'none', the scores for each
+            class are returned. If 'macro', calculate metrics for each class,
+            and find their unweighted mean.
+            Defaults to 'macro'.
+        thrs (float | tuple[float], optional): Predictions with scores under
+            the thresholds are considered negative. Default to None.
+
+    Returns:
+         float | np.array | list[float | np.array]: Precision.
+            If the ``average_mode`` is set to macro, np.array is used in favor
+            of float to give class-wise results. If the ``average_mode`` is set
+             to none, float is used to return a single value.
+            If ``thrs`` is a single float or None, the function will return
+            float or np.array. If ``thrs`` is a tuple, the function will return
+             a list containing metrics for each ``thrs`` condition.
+    """
+    precisions, _, _ = precision_recall_f1(pred, target, average_mode, thrs)
+    return precisions
+
+
+def recall(pred, target, average_mode='macro', thrs=None):
+    """Calculate recall according to the prediction and target.
+
+    Args:
+        pred (torch.Tensor | np.array): The model prediction with shape (N, C).
+        target (torch.Tensor | np.array): The target of each prediction with
+            shape (N, 1) or (N,).
+        average_mode (str): The type of averaging performed on the result.
+            Options are 'macro' and 'none'. If 'none', the scores for each
+            class are returned. If 'macro', calculate metrics for each class,
+            and find their unweighted mean.
+            Defaults to 'macro'.
+        thrs (float | tuple[float], optional): Predictions with scores under
+            the thresholds are considered negative. Default to None.
+
+    Returns:
+         float | np.array | list[float | np.array]: Recall.
+            If the ``average_mode`` is set to macro, np.array is used in favor
+            of float to give class-wise results. If the ``average_mode`` is set
+             to none, float is used to return a single value.
+            If ``thrs`` is a single float or None, the function will return
+            float or np.array. If ``thrs`` is a tuple, the function will return
+             a list containing metrics for each ``thrs`` condition.
+    """
+    _, recalls, _ = precision_recall_f1(pred, target, average_mode, thrs)
+    return recalls
+
+
+def f1_score(pred, target, average_mode='macro', thrs=None):
+    """Calculate F1 score according to the prediction and target.
+
+    Args:
+        pred (torch.Tensor | np.array): The model prediction with shape (N, C).
+        target (torch.Tensor | np.array): The target of each prediction with
+            shape (N, 1) or (N,).
+        average_mode (str): The type of averaging performed on the result.
+            Options are 'macro' and 'none'. If 'none', the scores for each
+            class are returned. If 'macro', calculate metrics for each class,
+            and find their unweighted mean.
+            Defaults to 'macro'.
+        thrs (float | tuple[float], optional): Predictions with scores under
+            the thresholds are considered negative. Default to None.
+
+    Returns:
+         float | np.array | list[float | np.array]: F1 score.
+            If the ``average_mode`` is set to macro, np.array is used in favor
+            of float to give class-wise results. If the ``average_mode`` is set
+             to none, float is used to return a single value.
+            If ``thrs`` is a single float or None, the function will return
+            float or np.array. If ``thrs`` is a tuple, the function will return
+             a list containing metrics for each ``thrs`` condition.
+    """
+    _, _, f1_scores = precision_recall_f1(pred, target, average_mode, thrs)
+    return f1_scores
+
+
+def support(pred, target, average_mode='macro'):
+    """Calculate the total number of occurrences of each label according to the
+    prediction and target.
+
+    Args:
+        pred (torch.Tensor | np.array): The model prediction with shape (N, C).
+        target (torch.Tensor | np.array): The target of each prediction with
+            shape (N, 1) or (N,).
+        average_mode (str): The type of averaging performed on the result.
+            Options are 'macro' and 'none'. If 'none', the scores for each
+            class are returned. If 'macro', calculate metrics for each class,
+            and find their unweighted sum.
+            Defaults to 'macro'.
+
+    Returns:
+        float | np.array: Precision, recall, f1 score.
+            The function returns a single float if the average_mode is set to
+            macro, or a np.array with shape C if the average_mode is set to
+             none.
+    """
+    confusion_matrix = calculate_confusion_matrix(pred, target)
+    with torch.no_grad():
+        res = confusion_matrix.sum(1)
+        if average_mode == 'macro':
+            res = float(res.sum().numpy())
+        elif average_mode == 'none':
+            res = res.numpy()
+        else:
+            raise ValueError(f'Unsupport type of averaging {average_mode}.')
+    return res

openmmlab_test/mmclassification-speed-benchmark/mmcls/core/evaluation/mean_ap.py

+import numpy as np
+import torch
+
+
+def average_precision(pred, target):
+    """Calculate the average precision for a single class.
+
+    AP summarizes a precision-recall curve as the weighted mean of maximum
+    precisions obtained for any r'>r, where r is the recall:
+
+    ..math::
+        \\text{AP} = \\sum_n (R_n - R_{n-1}) P_n
+
+    Note that no approximation is involved since the curve is piecewise
+    constant.
+
+    Args:
+        pred (np.ndarray): The model prediction with shape (N, ).
+        target (np.ndarray): The target of each prediction with shape (N, ).
+
+    Returns:
+        float: a single float as average precision value.
+    """
+    eps = np.finfo(np.float32).eps
+
+    # sort examples
+    sort_inds = np.argsort(-pred)
+    sort_target = target[sort_inds]
+
+    # count true positive examples
+    pos_inds = sort_target == 1
+    tp = np.cumsum(pos_inds)
+    total_pos = tp[-1]
+
+    # count not difficult examples
+    pn_inds = sort_target != -1
+    pn = np.cumsum(pn_inds)
+
+    tp[np.logical_not(pos_inds)] = 0
+    precision = tp / np.maximum(pn, eps)
+    ap = np.sum(precision) / np.maximum(total_pos, eps)
+    return ap
+
+
+def mAP(pred, target):
+    """Calculate the mean average precision with respect of classes.
+
+    Args:
+        pred (torch.Tensor | np.ndarray): The model prediction with shape
+            (N, C), where C is the number of classes.
+        target (torch.Tensor | np.ndarray): The target of each prediction with
+            shape (N, C), where C is the number of classes. 1 stands for
+            positive examples, 0 stands for negative examples and -1 stands for
+            difficult examples.
+
+    Returns:
+        float: A single float as mAP value.
+    """
+    if isinstance(pred, torch.Tensor) and isinstance(target, torch.Tensor):
+        pred = pred.detach().cpu().numpy()
+        target = target.detach().cpu().numpy()
+    elif not (isinstance(pred, np.ndarray) and isinstance(target, np.ndarray)):
+        raise TypeError('pred and target should both be torch.Tensor or'
+                        'np.ndarray')
+
+    assert pred.shape == \
+        target.shape, 'pred and target should be in the same shape.'
+    num_classes = pred.shape[1]
+    ap = np.zeros(num_classes)
+    for k in range(num_classes):
+        ap[k] = average_precision(pred[:, k], target[:, k])
+    mean_ap = ap.mean() * 100.0
+    return mean_ap

openmmlab_test/mmclassification-speed-benchmark/mmcls/core/evaluation/multilabel_eval_metrics.py

+import warnings
+
+import numpy as np
+import torch
+
+
+def average_performance(pred, target, thr=None, k=None):
+    """Calculate CP, CR, CF1, OP, OR, OF1, where C stands for per-class
+    average, O stands for overall average, P stands for precision, R stands for
+    recall and F1 stands for F1-score.
+
+    Args:
+        pred (torch.Tensor | np.ndarray): The model prediction with shape
+            (N, C), where C is the number of classes.
+        target (torch.Tensor | np.ndarray): The target of each prediction with
+            shape (N, C), where C is the number of classes. 1 stands for
+            positive examples, 0 stands for negative examples and -1 stands for
+            difficult examples.
+        thr (float): The confidence threshold. Defaults to None.
+        k (int): Top-k performance. Note that if thr and k are both given, k
+            will be ignored. Defaults to None.
+
+    Returns:
+        tuple: (CP, CR, CF1, OP, OR, OF1)
+    """
+    if isinstance(pred, torch.Tensor) and isinstance(target, torch.Tensor):
+        pred = pred.detach().cpu().numpy()
+        target = target.detach().cpu().numpy()
+    elif not (isinstance(pred, np.ndarray) and isinstance(target, np.ndarray)):
+        raise TypeError('pred and target should both be torch.Tensor or'
+                        'np.ndarray')
+    if thr is None and k is None:
+        thr = 0.5
+        warnings.warn('Neither thr nor k is given, set thr as 0.5 by '
+                      'default.')
+    elif thr is not None and k is not None:
+        warnings.warn('Both thr and k are given, use threshold in favor of '
+                      'top-k.')
+
+    assert pred.shape == \
+        target.shape, 'pred and target should be in the same shape.'
+
+    eps = np.finfo(np.float32).eps
+    target[target == -1] = 0
+    if thr is not None:
+        # a label is predicted positive if the confidence is no lower than thr
+        pos_inds = pred >= thr
+
+    else:
+        # top-k labels will be predicted positive for any example
+        sort_inds = np.argsort(-pred, axis=1)
+        sort_inds_ = sort_inds[:, :k]
+        inds = np.indices(sort_inds_.shape)
+        pos_inds = np.zeros_like(pred)
+        pos_inds[inds[0], sort_inds_] = 1
+
+    tp = (pos_inds * target) == 1
+    fp = (pos_inds * (1 - target)) == 1
+    fn = ((1 - pos_inds) * target) == 1
+
+    precision_class = tp.sum(axis=0) / np.maximum(
+        tp.sum(axis=0) + fp.sum(axis=0), eps)
+    recall_class = tp.sum(axis=0) / np.maximum(
+        tp.sum(axis=0) + fn.sum(axis=0), eps)
+    CP = precision_class.mean() * 100.0
+    CR = recall_class.mean() * 100.0
+    CF1 = 2 * CP * CR / np.maximum(CP + CR, eps)
+    OP = tp.sum() / np.maximum(tp.sum() + fp.sum(), eps) * 100.0
+    OR = tp.sum() / np.maximum(tp.sum() + fn.sum(), eps) * 100.0
+    OF1 = 2 * OP * OR / np.maximum(OP + OR, eps)
+    return CP, CR, CF1, OP, OR, OF1

openmmlab_test/mmclassification-speed-benchmark/mmcls/core/export/__init__.py

+from .test import ONNXRuntimeClassifier, TensorRTClassifier
+
+__all__ = ['ONNXRuntimeClassifier', 'TensorRTClassifier']

openmmlab_test/mmclassification-speed-benchmark/mmcls/core/export/test.py

+import warnings
+
+import numpy as np
+import onnxruntime as ort
+import torch
+
+from mmcls.models.classifiers import BaseClassifier
+
+
+class ONNXRuntimeClassifier(BaseClassifier):
+    """Wrapper for classifier's inference with ONNXRuntime."""
+
+    def __init__(self, onnx_file, class_names, device_id):
+        super(ONNXRuntimeClassifier, self).__init__()
+        sess = ort.InferenceSession(onnx_file)
+
+        providers = ['CPUExecutionProvider']
+        options = [{}]
+        is_cuda_available = ort.get_device() == 'GPU'
+        if is_cuda_available:
+            providers.insert(0, 'CUDAExecutionProvider')
+            options.insert(0, {'device_id': device_id})
+        sess.set_providers(providers, options)
+
+        self.sess = sess
+        self.CLASSES = class_names
+        self.device_id = device_id
+        self.io_binding = sess.io_binding()
+        self.output_names = [_.name for _ in sess.get_outputs()]
+        self.is_cuda_available = is_cuda_available
+
+    def simple_test(self, img, img_metas, **kwargs):
+        raise NotImplementedError('This method is not implemented.')
+
+    def extract_feat(self, imgs):
+        raise NotImplementedError('This method is not implemented.')
+
+    def forward_train(self, imgs, **kwargs):
+        raise NotImplementedError('This method is not implemented.')
+
+    def forward_test(self, imgs, img_metas, **kwargs):
+        input_data = imgs
+        # set io binding for inputs/outputs
+        device_type = 'cuda' if self.is_cuda_available else 'cpu'
+        if not self.is_cuda_available:
+            input_data = input_data.cpu()
+        self.io_binding.bind_input(
+            name='input',
+            device_type=device_type,
+            device_id=self.device_id,
+            element_type=np.float32,
+            shape=input_data.shape,
+            buffer_ptr=input_data.data_ptr())
+
+        for name in self.output_names:
+            self.io_binding.bind_output(name)
+        # run session to get outputs
+        self.sess.run_with_iobinding(self.io_binding)
+        results = self.io_binding.copy_outputs_to_cpu()[0]
+        return list(results)
+
+
+class TensorRTClassifier(BaseClassifier):
+
+    def __init__(self, trt_file, class_names, device_id):
+        super(TensorRTClassifier, self).__init__()
+        from mmcv.tensorrt import TRTWraper, load_tensorrt_plugin
+        try:
+            load_tensorrt_plugin()
+        except (ImportError, ModuleNotFoundError):
+            warnings.warn('If input model has custom op from mmcv, \
+                you may have to build mmcv with TensorRT from source.')
+        model = TRTWraper(
+            trt_file, input_names=['input'], output_names=['probs'])
+
+        self.model = model
+        self.device_id = device_id
+        self.CLASSES = class_names
+
+    def simple_test(self, img, img_metas, **kwargs):
+        raise NotImplementedError('This method is not implemented.')
+
+    def extract_feat(self, imgs):
+        raise NotImplementedError('This method is not implemented.')
+
+    def forward_train(self, imgs, **kwargs):
+        raise NotImplementedError('This method is not implemented.')
+
+    def forward_test(self, imgs, img_metas, **kwargs):
+        input_data = imgs
+        with torch.cuda.device(self.device_id), torch.no_grad():
+            results = self.model({'input': input_data})['probs']
+        results = results.detach().cpu().numpy()
+
+        return list(results)

openmmlab_test/mmclassification-speed-benchmark/mmcls/core/fp16/__init__.py

+from .decorators import auto_fp16, force_fp32
+from .hooks import Fp16OptimizerHook, wrap_fp16_model
+
+__all__ = ['auto_fp16', 'force_fp32', 'Fp16OptimizerHook', 'wrap_fp16_model']

openmmlab_test/mmclassification-speed-benchmark/mmcls/core/fp16/decorators.py

+import functools
+from inspect import getfullargspec
+
+import torch
+
+from .utils import cast_tensor_type
+
+
+def auto_fp16(apply_to=None, out_fp32=False):
+    """Decorator to enable fp16 training automatically.
+
+    This decorator is useful when you write custom modules and want to support
+    mixed precision training. If inputs arguments are fp32 tensors, they will
+    be converted to fp16 automatically. Arguments other than fp32 tensors are
+    ignored.
+
+    Args:
+        apply_to (Iterable, optional): The argument names to be converted.
+            `None` indicates all arguments.
+        out_fp32 (bool): Whether to convert the output back to fp32.
+
+    :Example:
+
+        class MyModule1(nn.Module)
+
+            # Convert x and y to fp16
+            @auto_fp16()
+            def forward(self, x, y):
+                pass
+
+        class MyModule2(nn.Module):
+
+            # convert pred to fp16
+            @auto_fp16(apply_to=('pred', ))
+            def do_something(self, pred, others):
+                pass
+    """
+
+    def auto_fp16_wrapper(old_func):
+
+        @functools.wraps(old_func)
+        def new_func(*args, **kwargs):
+            # check if the module has set the attribute `fp16_enabled`, if not,
+            # just fallback to the original method.
+            if not isinstance(args[0], torch.nn.Module):
+                raise TypeError('@auto_fp16 can only be used to decorate the '
+                                'method of nn.Module')
+            if not (hasattr(args[0], 'fp16_enabled') and args[0].fp16_enabled):
+                return old_func(*args, **kwargs)
+            # get the arg spec of the decorated method
+            args_info = getfullargspec(old_func)
+            # get the argument names to be casted
+            args_to_cast = args_info.args if apply_to is None else apply_to
+            # convert the args that need to be processed
+            new_args = []
+            # NOTE: default args are not taken into consideration
+            if args:
+                arg_names = args_info.args[:len(args)]
+                for i, arg_name in enumerate(arg_names):
+                    if arg_name in args_to_cast:
+                        new_args.append(
+                            cast_tensor_type(args[i], torch.float, torch.half))
+                    else:
+                        new_args.append(args[i])
+            # convert the kwargs that need to be processed
+            new_kwargs = {}
+            if kwargs:
+                for arg_name, arg_value in kwargs.items():
+                    if arg_name in args_to_cast:
+                        new_kwargs[arg_name] = cast_tensor_type(
+                            arg_value, torch.float, torch.half)
+                    else:
+                        new_kwargs[arg_name] = arg_value
+            # apply converted arguments to the decorated method
+            output = old_func(*new_args, **new_kwargs)
+            # cast the results back to fp32 if necessary
+            if out_fp32:
+                output = cast_tensor_type(output, torch.half, torch.float)
+            return output
+
+        return new_func
+
+    return auto_fp16_wrapper
+
+
+def force_fp32(apply_to=None, out_fp16=False):
+    """Decorator to convert input arguments to fp32 in force.
+
+    This decorator is useful when you write custom modules and want to support
+    mixed precision training. If there are some inputs that must be processed
+    in fp32 mode, then this decorator can handle it. If inputs arguments are
+    fp16 tensors, they will be converted to fp32 automatically. Arguments other
+    than fp16 tensors are ignored.
+
+    Args:
+        apply_to (Iterable, optional): The argument names to be converted.
+            `None` indicates all arguments.
+        out_fp16 (bool): Whether to convert the output back to fp16.
+
+    :Example:
+
+        class MyModule1(nn.Module)
+
+            # Convert x and y to fp32
+            @force_fp32()
+            def loss(self, x, y):
+                pass
+
+        class MyModule2(nn.Module):
+
+            # convert pred to fp32
+            @force_fp32(apply_to=('pred', ))
+            def post_process(self, pred, others):
+                pass
+    """
+
+    def force_fp32_wrapper(old_func):
+
+        @functools.wraps(old_func)
+        def new_func(*args, **kwargs):
+            # check if the module has set the attribute `fp16_enabled`, if not,
+            # just fallback to the original method.
+            if not isinstance(args[0], torch.nn.Module):
+                raise TypeError('@force_fp32 can only be used to decorate the '
+                                'method of nn.Module')
+            if not (hasattr(args[0], 'fp16_enabled') and args[0].fp16_enabled):
+                return old_func(*args, **kwargs)
+            # get the arg spec of the decorated method
+            args_info = getfullargspec(old_func)
+            # get the argument names to be casted
+            args_to_cast = args_info.args if apply_to is None else apply_to
+            # convert the args that need to be processed
+            new_args = []
+            if args:
+                arg_names = args_info.args[:len(args)]
+                for i, arg_name in enumerate(arg_names):
+                    if arg_name in args_to_cast:
+                        new_args.append(
+                            cast_tensor_type(args[i], torch.half, torch.float))
+                    else:
+                        new_args.append(args[i])
+            # convert the kwargs that need to be processed
+            new_kwargs = dict()
+            if kwargs:
+                for arg_name, arg_value in kwargs.items():
+                    if arg_name in args_to_cast:
+                        new_kwargs[arg_name] = cast_tensor_type(
+                            arg_value, torch.half, torch.float)
+                    else:
+                        new_kwargs[arg_name] = arg_value
+            # apply converted arguments to the decorated method
+            output = old_func(*new_args, **new_kwargs)
+            # cast the results back to fp32 if necessary
+            if out_fp16:
+                output = cast_tensor_type(output, torch.float, torch.half)
+            return output
+
+        return new_func
+
+    return force_fp32_wrapper

openmmlab_test/mmclassification-speed-benchmark/mmcls/core/fp16/hooks.py

+import copy
+
+import torch
+import torch.nn as nn
+from mmcv.runner import OptimizerHook
+from mmcv.utils.parrots_wrapper import _BatchNorm
+
+from ..utils import allreduce_grads
+from .utils import cast_tensor_type
+
+
+class Fp16OptimizerHook(OptimizerHook):
+    """FP16 optimizer hook.
+
+    The steps of fp16 optimizer is as follows.
+    1. Scale the loss value.
+    2. BP in the fp16 model.
+    2. Copy gradients from fp16 model to fp32 weights.
+    3. Update fp32 weights.
+    4. Copy updated parameters from fp32 weights to fp16 model.
+
+    Refer to https://arxiv.org/abs/1710.03740 for more details.
+
+    Args:
+        loss_scale (float): Scale factor multiplied with loss.
+    """
+
+    def __init__(self,
+                 grad_clip=None,
+                 coalesce=True,
+                 bucket_size_mb=-1,
+                 loss_scale=512.,
+                 distributed=True):
+        self.grad_clip = grad_clip
+        self.coalesce = coalesce
+        self.bucket_size_mb = bucket_size_mb
+        self.loss_scale = loss_scale
+        self.distributed = distributed
+
+    def before_run(self, runner):
+        # keep a copy of fp32 weights
+        runner.optimizer.param_groups = copy.deepcopy(
+            runner.optimizer.param_groups)
+        # convert model to fp16
+        wrap_fp16_model(runner.model)
+
+    def copy_grads_to_fp32(self, fp16_net, fp32_weights):
+        """Copy gradients from fp16 model to fp32 weight copy."""
+        for fp32_param, fp16_param in zip(fp32_weights, fp16_net.parameters()):
+            if fp16_param.grad is not None:
+                if fp32_param.grad is None:
+                    fp32_param.grad = fp32_param.data.new(fp32_param.size())
+                fp32_param.grad.copy_(fp16_param.grad)
+
+    def copy_params_to_fp16(self, fp16_net, fp32_weights):
+        """Copy updated params from fp32 weight copy to fp16 model."""
+        for fp16_param, fp32_param in zip(fp16_net.parameters(), fp32_weights):
+            fp16_param.data.copy_(fp32_param.data)
+
+    def after_train_iter(self, runner):
+        # clear grads of last iteration
+        runner.model.zero_grad()
+        runner.optimizer.zero_grad()
+        # scale the loss value
+        scaled_loss = runner.outputs['loss'] * self.loss_scale
+        scaled_loss.backward()
+        # copy fp16 grads in the model to fp32 params in the optimizer
+        fp32_weights = []
+        for param_group in runner.optimizer.param_groups:
+            fp32_weights += param_group['params']
+        self.copy_grads_to_fp32(runner.model, fp32_weights)
+        # allreduce grads
+        if self.distributed:
+            allreduce_grads(fp32_weights, self.coalesce, self.bucket_size_mb)
+        # scale the gradients back
+        for param in fp32_weights:
+            if param.grad is not None:
+                param.grad.div_(self.loss_scale)
+        if self.grad_clip is not None:
+            self.clip_grads(fp32_weights)
+        # update fp32 params
+        runner.optimizer.step()
+        # copy fp32 params to the fp16 model
+        self.copy_params_to_fp16(runner.model, fp32_weights)
+
+
+def wrap_fp16_model(model):
+    # convert model to fp16
+    model.half()
+    # patch the normalization layers to make it work in fp32 mode
+    patch_norm_fp32(model)
+    # set `fp16_enabled` flag
+    for m in model.modules():
+        if hasattr(m, 'fp16_enabled'):
+            m.fp16_enabled = True
+
+
+def patch_norm_fp32(module):
+    if isinstance(module, (_BatchNorm, nn.GroupNorm)):
+        module.float()
+        module.forward = patch_forward_method(module.forward, torch.half,
+                                              torch.float)
+    for child in module.children():
+        patch_norm_fp32(child)
+    return module
+
+
+def patch_forward_method(func, src_type, dst_type, convert_output=True):
+    """Patch the forward method of a module.
+
+    Args:
+        func (callable): The original forward method.
+        src_type (torch.dtype): Type of input arguments to be converted from.
+        dst_type (torch.dtype): Type of input arguments to be converted to.
+        convert_output (bool): Whether to convert the output back to src_type.
+
+    Returns:
+        callable: The patched forward method.
+    """
+
+    def new_forward(*args, **kwargs):
+        output = func(*cast_tensor_type(args, src_type, dst_type),
+                      **cast_tensor_type(kwargs, src_type, dst_type))
+        if convert_output:
+            output = cast_tensor_type(output, dst_type, src_type)
+        return output
+
+    return new_forward

openmmlab_test/mmclassification-speed-benchmark/mmcls/core/fp16/utils.py

+from collections import abc
+
+import numpy as np
+import torch
+
+
+def cast_tensor_type(inputs, src_type, dst_type):
+    if isinstance(inputs, torch.Tensor):
+        return inputs.to(dst_type)
+    elif isinstance(inputs, str):
+        return inputs
+    elif isinstance(inputs, np.ndarray):
+        return inputs
+    elif isinstance(inputs, abc.Mapping):
+        return type(inputs)({
+            k: cast_tensor_type(v, src_type, dst_type)
+            for k, v in inputs.items()
+        })
+    elif isinstance(inputs, abc.Iterable):
+        return type(inputs)(
+            cast_tensor_type(item, src_type, dst_type) for item in inputs)
+    else:
+        return inputs

openmmlab_test/mmclassification-speed-benchmark/mmcls/core/utils/__init__.py

+from .dist_utils import DistOptimizerHook, allreduce_grads
+from .misc import multi_apply
+
+__all__ = ['allreduce_grads', 'DistOptimizerHook', 'multi_apply']