We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[bug] error when using `python tools/train.py` to train `ReDet` in MMRotate.
when the `self.distributed` equals to False, `revert_sync_batchnorm` function cannot work normally. 
https://github.com/open-mmlab/mmengine/blob/bda92b49e751c721c11bb15623dabc7f7a9517e7/mmengine/runner/runner.py#L848

We argue this is caused by [ReDet](https://github.com/open-mmlab/mmrotate/blob/dev-1.x/configs/redet/redet_re50_refpn_1x_dota_le90.py) replacing ``torch.nn`` with ``e2cnn.nn``.

https://github.com/open-mmlab/mmrotate/blob/dev-1.x/mmrotate/models/backbones/re_resnet.py



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mmengine/model/utils.py`
Content:
```
1 # Copyright (c) OpenMMLab. All rights reserved.
2 import logging
3 import warnings
4 from typing import List, Union
5 
6 import torch
7 import torch.nn as nn
8 import torch.nn.functional as F
9 
10 from mmengine.utils.dl_utils import mmcv_full_available
11 
12 
13 def stack_batch(tensor_list: List[torch.Tensor],
14                 pad_size_divisor: int = 1,
15                 pad_value: Union[int, float] = 0) -> torch.Tensor:
16     """Stack multiple tensors to form a batch and pad the tensor to the max
17     shape use the right bottom padding mode in these images. If
18     ``pad_size_divisor > 0``, add padding to ensure the shape of each dim is
19     divisible by ``pad_size_divisor``.
20 
21     Args:
22         tensor_list (List[Tensor]): A list of tensors with the same dim.
23         pad_size_divisor (int): If ``pad_size_divisor > 0``, add padding
24             to ensure the shape of each dim is divisible by
25             ``pad_size_divisor``. This depends on the model, and many
26             models need to be divisible by 32. Defaults to 1
27         pad_value (int, float): The padding value. Defaults to 0.
28 
29     Returns:
30        Tensor: The n dim tensor.
31     """
32     assert isinstance(
33         tensor_list,
34         list), (f'Expected input type to be list, but got {type(tensor_list)}')
35     assert tensor_list, '`tensor_list` could not be an empty list'
36     assert len({
37         tensor.ndim
38         for tensor in tensor_list
39     }) == 1, (f'Expected the dimensions of all tensors must be the same, '
40               f'but got {[tensor.ndim for tensor in tensor_list]}')
41 
42     dim = tensor_list[0].dim()
43     num_img = len(tensor_list)
44     all_sizes: torch.Tensor = torch.Tensor(
45         [tensor.shape for tensor in tensor_list])
46     max_sizes = torch.ceil(
47         torch.max(all_sizes, dim=0)[0] / pad_size_divisor) * pad_size_divisor
48     padded_sizes = max_sizes - all_sizes
49     # The first dim normally means channel,  which should not be padded.
50     padded_sizes[:, 0] = 0
51     if padded_sizes.sum() == 0:
52         return torch.stack(tensor_list)
53     # `pad` is the second arguments of `F.pad`. If pad is (1, 2, 3, 4),
54     # it means that padding the last dim with 1(left) 2(right), padding the
55     # penultimate dim to 3(top) 4(bottom). The order of `pad` is opposite of
56     # the `padded_sizes`. Therefore, the `padded_sizes` needs to be reversed,
57     # and only odd index of pad should be assigned to keep padding "right" and
58     # "bottom".
59     pad = torch.zeros(num_img, 2 * dim, dtype=torch.int)
60     pad[:, 1::2] = padded_sizes[:, range(dim - 1, -1, -1)]
61     batch_tensor = []
62     for idx, tensor in enumerate(tensor_list):
63         batch_tensor.append(
64             F.pad(tensor, tuple(pad[idx].tolist()), value=pad_value))
65     return torch.stack(batch_tensor)
66 
67 
68 def detect_anomalous_params(loss: torch.Tensor, model) -> None:
69     parameters_in_graph = set()
70     visited = set()
71 
72     def traverse(grad_fn):
73         if grad_fn is None:
74             return
75         if grad_fn not in visited:
76             visited.add(grad_fn)
77             if hasattr(grad_fn, 'variable'):
78                 parameters_in_graph.add(grad_fn.variable)
79             parents = grad_fn.next_functions
80             if parents is not None:
81                 for parent in parents:
82                     grad_fn = parent[0]
83                     traverse(grad_fn)
84 
85     traverse(loss.grad_fn)
86     from mmengine.logging import MMLogger
87     logger = MMLogger.get_current_instance()
88     for n, p in model.named_parameters():
89         if p not in parameters_in_graph and p.requires_grad:
90             logger.log(
91                 level=logging.ERROR,
92                 msg=f'{n} with shape {p.size()} is not '
93                 f'in the computational graph \n')
94 
95 
96 def merge_dict(*args):
97     """Merge all dictionaries into one dictionary.
98 
99     If pytorch version >= 1.8, ``merge_dict`` will be wrapped
100     by ``torch.fx.wrap``,  which will make ``torch.fx.symbolic_trace`` skip
101     trace ``merge_dict``.
102 
103     Note:
104         If a function needs to be traced by ``torch.fx.symbolic_trace``,
105         but inevitably needs to use ``update`` method of ``dict``(``update``
106         is not traceable). It should use ``merge_dict`` to replace
107         ``xxx.update``.
108 
109     Args:
110         *args: dictionary needs to be merged.
111 
112     Returns:
113         dict: Merged dict from args
114     """
115     output = dict()
116     for item in args:
117         assert isinstance(
118             item,
119             dict), (f'all arguments of merge_dict should be a dict, but got '
120                     f'{type(item)}')
121         output.update(item)
122     return output
123 
124 
125 # torch.fx is only available when pytorch version >= 1.8.
126 # If the subclass of `BaseModel` has multiple submodules, and each module
127 # will return a loss dict during training process, i.e., `TwoStageDetector`
128 # in mmdet. It should use `merge_dict` to get the total loss, rather than
129 # `loss.update` to keep model traceable.
130 try:
131     import torch.fx
132 
133     # make torch.fx skip trace `merge_dict`.
134     merge_dict = torch.fx.wrap(merge_dict)
135 
136 except ImportError:
137     warnings.warn('Cannot import torch.fx, `merge_dict` is a simple function '
138                   'to merge multiple dicts')
139 
140 
141 class _BatchNormXd(nn.modules.batchnorm._BatchNorm):
142     """A general BatchNorm layer without input dimension check.
143 
144     Reproduced from @kapily's work:
145     (https://github.com/pytorch/pytorch/issues/41081#issuecomment-783961547)
146     The only difference between BatchNorm1d, BatchNorm2d, BatchNorm3d, etc
147     is `_check_input_dim` that is designed for tensor sanity checks.
148     The check has been bypassed in this class for the convenience of converting
149     SyncBatchNorm.
150     """
151 
152     def _check_input_dim(self, input: torch.Tensor):
153         return
154 
155 
156 def revert_sync_batchnorm(module: nn.Module) -> nn.Module:
157     """Helper function to convert all `SyncBatchNorm` (SyncBN) and
158     `mmcv.ops.sync_bn.SyncBatchNorm`(MMSyncBN) layers in the model to
159     `BatchNormXd` layers.
160 
161     Adapted from @kapily's work:
162     (https://github.com/pytorch/pytorch/issues/41081#issuecomment-783961547)
163 
164     Args:
165         module (nn.Module): The module containing `SyncBatchNorm` layers.
166 
167     Returns:
168         module_output: The converted module with `BatchNormXd` layers.
169     """
170     module_output = module
171     module_checklist = [torch.nn.modules.batchnorm.SyncBatchNorm]
172 
173     if mmcv_full_available():
174         from mmcv.ops import SyncBatchNorm
175         module_checklist.append(SyncBatchNorm)
176 
177     if isinstance(module, tuple(module_checklist)):
178         module_output = _BatchNormXd(module.num_features, module.eps,
179                                      module.momentum, module.affine,
180                                      module.track_running_stats)
181         if module.affine:
182             # no_grad() may not be needed here but
183             # just to be consistent with `convert_sync_batchnorm()`
184             with torch.no_grad():
185                 module_output.weight = module.weight
186                 module_output.bias = module.bias
187         module_output.running_mean = module.running_mean
188         module_output.running_var = module.running_var
189         module_output.num_batches_tracked = module.num_batches_tracked
190         module_output.training = module.training
191         # qconfig exists in quantized models
192         if hasattr(module, 'qconfig'):
193             module_output.qconfig = module.qconfig
194     for name, child in module.named_children():
195         module_output.add_module(name, revert_sync_batchnorm(child))
196     del module
197     return module_output
198 
199 
200 def convert_sync_batchnorm(module: nn.Module,
201                            implementation='torch') -> nn.Module:
202     """Helper function to convert all `BatchNorm` layers in the model to
203     `SyncBatchNorm` (SyncBN) or `mmcv.ops.sync_bn.SyncBatchNorm`(MMSyncBN)
204     layers. Adapted from <https://pytorch.org/docs/stable/generated/torch.nn.Sy
205     ncBatchNorm.html#torch.nn.SyncBatchNorm.convert_sync_batchnorm>_.
206 
207     Args:
208         module (nn.Module): The module containing `SyncBatchNorm` layers.
209         implementation (str): The type of `SyncBatchNorm` to convert to.
210 
211             - 'torch': convert to `torch.nn.modules.batchnorm.SyncBatchNorm`.
212             - 'mmcv': convert to `mmcv.ops.sync_bn.SyncBatchNorm`.
213 
214     Returns:
215         nn.Module: The converted module with `SyncBatchNorm` layers.
216     """  # noqa: E501
217     module_output = module
218 
219     if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
220         if implementation == 'torch':
221             SyncBatchNorm = torch.nn.modules.batchnorm.SyncBatchNorm
222         elif implementation == 'mmcv':
223             from mmcv.ops import SyncBatchNorm  # type: ignore
224         else:
225             raise ValueError('sync_bn should be "torch" or "mmcv", but got '
226                              f'{implementation}')
227 
228         module_output = SyncBatchNorm(module.num_features, module.eps,
229                                       module.momentum, module.affine,
230                                       module.track_running_stats)
231 
232         if module.affine:
233             with torch.no_grad():
234                 module_output.weight = module.weight
235                 module_output.bias = module.bias
236         module_output.running_mean = module.running_mean
237         module_output.running_var = module.running_var
238         module_output.num_batches_tracked = module.num_batches_tracked
239         if hasattr(module, 'qconfig'):
240             module_output.qconfig = module.qconfig
241     for name, child in module.named_children():
242         module_output.add_module(name,
243                                  convert_sync_batchnorm(child, implementation))
244     del module
245     return module_output
246
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/mmengine/model/utils.py b/mmengine/model/utils.py
--- a/mmengine/model/utils.py
+++ b/mmengine/model/utils.py
@@ -7,6 +7,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+from mmengine.logging import print_log
 from mmengine.utils.dl_utils import mmcv_full_available
 
 
@@ -192,7 +193,17 @@
         if hasattr(module, 'qconfig'):
             module_output.qconfig = module.qconfig
     for name, child in module.named_children():
-        module_output.add_module(name, revert_sync_batchnorm(child))
+        # Some custom modules or 3rd party implemented modules may raise an
+        # error when calling `add_module`. Therefore, try to catch the error
+        # and do not raise it. See https://github.com/open-mmlab/mmengine/issues/638 # noqa: E501
+        # for more details.
+        try:
+            module_output.add_module(name, revert_sync_batchnorm(child))
+        except Exception:
+            print_log(
+                F'Failed to convert {child} from SyncBN to BN!',
+                logger='current',
+                level=logging.WARNING)
     del module
     return module_output

{"golden_diff": "diff --git a/mmengine/model/utils.py b/mmengine/model/utils.py\n--- a/mmengine/model/utils.py\n+++ b/mmengine/model/utils.py\n@@ -7,6 +7,7 @@\n import torch.nn as nn\n import torch.nn.functional as F\n \n+from mmengine.logging import print_log\n from mmengine.utils.dl_utils import mmcv_full_available\n \n \n@@ -192,7 +193,17 @@\n         if hasattr(module, 'qconfig'):\n             module_output.qconfig = module.qconfig\n     for name, child in module.named_children():\n-        module_output.add_module(name, revert_sync_batchnorm(child))\n+        # Some custom modules or 3rd party implemented modules may raise an\n+        # error when calling `add_module`. Therefore, try to catch the error\n+        # and do not raise it. See https://github.com/open-mmlab/mmengine/issues/638 # noqa: E501\n+        # for more details.\n+        try:\n+            module_output.add_module(name, revert_sync_batchnorm(child))\n+        except Exception:\n+            print_log(\n+                F'Failed to convert {child} from SyncBN to BN!',\n+                logger='current',\n+                level=logging.WARNING)\n     del module\n     return module_output\n", "issue": "[bug] error when using `python tools/train.py` to train `ReDet` in MMRotate.\nwhen the `self.distributed` equals to False, `revert_sync_batchnorm` function cannot work normally. \r\nhttps://github.com/open-mmlab/mmengine/blob/bda92b49e751c721c11bb15623dabc7f7a9517e7/mmengine/runner/runner.py#L848\r\n\r\nWe argue this is caused by [ReDet](https://github.com/open-mmlab/mmrotate/blob/dev-1.x/configs/redet/redet_re50_refpn_1x_dota_le90.py) replacing ``torch.nn`` with ``e2cnn.nn``.\r\n\r\nhttps://github.com/open-mmlab/mmrotate/blob/dev-1.x/mmrotate/models/backbones/re_resnet.py\r\n\r\n\n", "before_files": [{"content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport logging\nimport warnings\nfrom typing import List, Union\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom mmengine.utils.dl_utils import mmcv_full_available\n\n\ndef stack_batch(tensor_list: List[torch.Tensor],\n                pad_size_divisor: int = 1,\n                pad_value: Union[int, float] = 0) -> torch.Tensor:\n    \"\"\"Stack multiple tensors to form a batch and pad the tensor to the max\n    shape use the right bottom padding mode in these images. If\n    ``pad_size_divisor > 0``, add padding to ensure the shape of each dim is\n    divisible by ``pad_size_divisor``.\n\n    Args:\n        tensor_list (List[Tensor]): A list of tensors with the same dim.\n        pad_size_divisor (int): If ``pad_size_divisor > 0``, add padding\n            to ensure the shape of each dim is divisible by\n            ``pad_size_divisor``. This depends on the model, and many\n            models need to be divisible by 32. Defaults to 1\n        pad_value (int, float): The padding value. Defaults to 0.\n\n    Returns:\n       Tensor: The n dim tensor.\n    \"\"\"\n    assert isinstance(\n        tensor_list,\n        list), (f'Expected input type to be list, but got {type(tensor_list)}')\n    assert tensor_list, '`tensor_list` could not be an empty list'\n    assert len({\n        tensor.ndim\n        for tensor in tensor_list\n    }) == 1, (f'Expected the dimensions of all tensors must be the same, '\n              f'but got {[tensor.ndim for tensor in tensor_list]}')\n\n    dim = tensor_list[0].dim()\n    num_img = len(tensor_list)\n    all_sizes: torch.Tensor = torch.Tensor(\n        [tensor.shape for tensor in tensor_list])\n    max_sizes = torch.ceil(\n        torch.max(all_sizes, dim=0)[0] / pad_size_divisor) * pad_size_divisor\n    padded_sizes = max_sizes - all_sizes\n    # The first dim normally means channel,  which should not be padded.\n    padded_sizes[:, 0] = 0\n    if padded_sizes.sum() == 0:\n        return torch.stack(tensor_list)\n    # `pad` is the second arguments of `F.pad`. If pad is (1, 2, 3, 4),\n    # it means that padding the last dim with 1(left) 2(right), padding the\n    # penultimate dim to 3(top) 4(bottom). The order of `pad` is opposite of\n    # the `padded_sizes`. Therefore, the `padded_sizes` needs to be reversed,\n    # and only odd index of pad should be assigned to keep padding \"right\" and\n    # \"bottom\".\n    pad = torch.zeros(num_img, 2 * dim, dtype=torch.int)\n    pad[:, 1::2] = padded_sizes[:, range(dim - 1, -1, -1)]\n    batch_tensor = []\n    for idx, tensor in enumerate(tensor_list):\n        batch_tensor.append(\n            F.pad(tensor, tuple(pad[idx].tolist()), value=pad_value))\n    return torch.stack(batch_tensor)\n\n\ndef detect_anomalous_params(loss: torch.Tensor, model) -> None:\n    parameters_in_graph = set()\n    visited = set()\n\n    def traverse(grad_fn):\n        if grad_fn is None:\n            return\n        if grad_fn not in visited:\n            visited.add(grad_fn)\n            if hasattr(grad_fn, 'variable'):\n                parameters_in_graph.add(grad_fn.variable)\n            parents = grad_fn.next_functions\n            if parents is not None:\n                for parent in parents:\n                    grad_fn = parent[0]\n                    traverse(grad_fn)\n\n    traverse(loss.grad_fn)\n    from mmengine.logging import MMLogger\n    logger = MMLogger.get_current_instance()\n    for n, p in model.named_parameters():\n        if p not in parameters_in_graph and p.requires_grad:\n            logger.log(\n                level=logging.ERROR,\n                msg=f'{n} with shape {p.size()} is not '\n                f'in the computational graph \\n')\n\n\ndef merge_dict(*args):\n    \"\"\"Merge all dictionaries into one dictionary.\n\n    If pytorch version >= 1.8, ``merge_dict`` will be wrapped\n    by ``torch.fx.wrap``,  which will make ``torch.fx.symbolic_trace`` skip\n    trace ``merge_dict``.\n\n    Note:\n        If a function needs to be traced by ``torch.fx.symbolic_trace``,\n        but inevitably needs to use ``update`` method of ``dict``(``update``\n        is not traceable). It should use ``merge_dict`` to replace\n        ``xxx.update``.\n\n    Args:\n        *args: dictionary needs to be merged.\n\n    Returns:\n        dict: Merged dict from args\n    \"\"\"\n    output = dict()\n    for item in args:\n        assert isinstance(\n            item,\n            dict), (f'all arguments of merge_dict should be a dict, but got '\n                    f'{type(item)}')\n        output.update(item)\n    return output\n\n\n# torch.fx is only available when pytorch version >= 1.8.\n# If the subclass of `BaseModel` has multiple submodules, and each module\n# will return a loss dict during training process, i.e., `TwoStageDetector`\n# in mmdet. It should use `merge_dict` to get the total loss, rather than\n# `loss.update` to keep model traceable.\ntry:\n    import torch.fx\n\n    # make torch.fx skip trace `merge_dict`.\n    merge_dict = torch.fx.wrap(merge_dict)\n\nexcept ImportError:\n    warnings.warn('Cannot import torch.fx, `merge_dict` is a simple function '\n                  'to merge multiple dicts')\n\n\nclass _BatchNormXd(nn.modules.batchnorm._BatchNorm):\n    \"\"\"A general BatchNorm layer without input dimension check.\n\n    Reproduced from @kapily's work:\n    (https://github.com/pytorch/pytorch/issues/41081#issuecomment-783961547)\n    The only difference between BatchNorm1d, BatchNorm2d, BatchNorm3d, etc\n    is `_check_input_dim` that is designed for tensor sanity checks.\n    The check has been bypassed in this class for the convenience of converting\n    SyncBatchNorm.\n    \"\"\"\n\n    def _check_input_dim(self, input: torch.Tensor):\n        return\n\n\ndef revert_sync_batchnorm(module: nn.Module) -> nn.Module:\n    \"\"\"Helper function to convert all `SyncBatchNorm` (SyncBN) and\n    `mmcv.ops.sync_bn.SyncBatchNorm`(MMSyncBN) layers in the model to\n    `BatchNormXd` layers.\n\n    Adapted from @kapily's work:\n    (https://github.com/pytorch/pytorch/issues/41081#issuecomment-783961547)\n\n    Args:\n        module (nn.Module): The module containing `SyncBatchNorm` layers.\n\n    Returns:\n        module_output: The converted module with `BatchNormXd` layers.\n    \"\"\"\n    module_output = module\n    module_checklist = [torch.nn.modules.batchnorm.SyncBatchNorm]\n\n    if mmcv_full_available():\n        from mmcv.ops import SyncBatchNorm\n        module_checklist.append(SyncBatchNorm)\n\n    if isinstance(module, tuple(module_checklist)):\n        module_output = _BatchNormXd(module.num_features, module.eps,\n                                     module.momentum, module.affine,\n                                     module.track_running_stats)\n        if module.affine:\n            # no_grad() may not be needed here but\n            # just to be consistent with `convert_sync_batchnorm()`\n            with torch.no_grad():\n                module_output.weight = module.weight\n                module_output.bias = module.bias\n        module_output.running_mean = module.running_mean\n        module_output.running_var = module.running_var\n        module_output.num_batches_tracked = module.num_batches_tracked\n        module_output.training = module.training\n        # qconfig exists in quantized models\n        if hasattr(module, 'qconfig'):\n            module_output.qconfig = module.qconfig\n    for name, child in module.named_children():\n        module_output.add_module(name, revert_sync_batchnorm(child))\n    del module\n    return module_output\n\n\ndef convert_sync_batchnorm(module: nn.Module,\n                           implementation='torch') -> nn.Module:\n    \"\"\"Helper function to convert all `BatchNorm` layers in the model to\n    `SyncBatchNorm` (SyncBN) or `mmcv.ops.sync_bn.SyncBatchNorm`(MMSyncBN)\n    layers. Adapted from <https://pytorch.org/docs/stable/generated/torch.nn.Sy\n    ncBatchNorm.html#torch.nn.SyncBatchNorm.convert_sync_batchnorm>_.\n\n    Args:\n        module (nn.Module): The module containing `SyncBatchNorm` layers.\n        implementation (str): The type of `SyncBatchNorm` to convert to.\n\n            - 'torch': convert to `torch.nn.modules.batchnorm.SyncBatchNorm`.\n            - 'mmcv': convert to `mmcv.ops.sync_bn.SyncBatchNorm`.\n\n    Returns:\n        nn.Module: The converted module with `SyncBatchNorm` layers.\n    \"\"\"  # noqa: E501\n    module_output = module\n\n    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):\n        if implementation == 'torch':\n            SyncBatchNorm = torch.nn.modules.batchnorm.SyncBatchNorm\n        elif implementation == 'mmcv':\n            from mmcv.ops import SyncBatchNorm  # type: ignore\n        else:\n            raise ValueError('sync_bn should be \"torch\" or \"mmcv\", but got '\n                             f'{implementation}')\n\n        module_output = SyncBatchNorm(module.num_features, module.eps,\n                                      module.momentum, module.affine,\n                                      module.track_running_stats)\n\n        if module.affine:\n            with torch.no_grad():\n                module_output.weight = module.weight\n                module_output.bias = module.bias\n        module_output.running_mean = module.running_mean\n        module_output.running_var = module.running_var\n        module_output.num_batches_tracked = module.num_batches_tracked\n        if hasattr(module, 'qconfig'):\n            module_output.qconfig = module.qconfig\n    for name, child in module.named_children():\n        module_output.add_module(name,\n                                 convert_sync_batchnorm(child, implementation))\n    del module\n    return module_output\n", "path": "mmengine/model/utils.py"}], "after_files": [{"content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport logging\nimport warnings\nfrom typing import List, Union\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom mmengine.logging import print_log\nfrom mmengine.utils.dl_utils import mmcv_full_available\n\n\ndef stack_batch(tensor_list: List[torch.Tensor],\n                pad_size_divisor: int = 1,\n                pad_value: Union[int, float] = 0) -> torch.Tensor:\n    \"\"\"Stack multiple tensors to form a batch and pad the tensor to the max\n    shape use the right bottom padding mode in these images. If\n    ``pad_size_divisor > 0``, add padding to ensure the shape of each dim is\n    divisible by ``pad_size_divisor``.\n\n    Args:\n        tensor_list (List[Tensor]): A list of tensors with the same dim.\n        pad_size_divisor (int): If ``pad_size_divisor > 0``, add padding\n            to ensure the shape of each dim is divisible by\n            ``pad_size_divisor``. This depends on the model, and many\n            models need to be divisible by 32. Defaults to 1\n        pad_value (int, float): The padding value. Defaults to 0.\n\n    Returns:\n       Tensor: The n dim tensor.\n    \"\"\"\n    assert isinstance(\n        tensor_list,\n        list), (f'Expected input type to be list, but got {type(tensor_list)}')\n    assert tensor_list, '`tensor_list` could not be an empty list'\n    assert len({\n        tensor.ndim\n        for tensor in tensor_list\n    }) == 1, (f'Expected the dimensions of all tensors must be the same, '\n              f'but got {[tensor.ndim for tensor in tensor_list]}')\n\n    dim = tensor_list[0].dim()\n    num_img = len(tensor_list)\n    all_sizes: torch.Tensor = torch.Tensor(\n        [tensor.shape for tensor in tensor_list])\n    max_sizes = torch.ceil(\n        torch.max(all_sizes, dim=0)[0] / pad_size_divisor) * pad_size_divisor\n    padded_sizes = max_sizes - all_sizes\n    # The first dim normally means channel,  which should not be padded.\n    padded_sizes[:, 0] = 0\n    if padded_sizes.sum() == 0:\n        return torch.stack(tensor_list)\n    # `pad` is the second arguments of `F.pad`. If pad is (1, 2, 3, 4),\n    # it means that padding the last dim with 1(left) 2(right), padding the\n    # penultimate dim to 3(top) 4(bottom). The order of `pad` is opposite of\n    # the `padded_sizes`. Therefore, the `padded_sizes` needs to be reversed,\n    # and only odd index of pad should be assigned to keep padding \"right\" and\n    # \"bottom\".\n    pad = torch.zeros(num_img, 2 * dim, dtype=torch.int)\n    pad[:, 1::2] = padded_sizes[:, range(dim - 1, -1, -1)]\n    batch_tensor = []\n    for idx, tensor in enumerate(tensor_list):\n        batch_tensor.append(\n            F.pad(tensor, tuple(pad[idx].tolist()), value=pad_value))\n    return torch.stack(batch_tensor)\n\n\ndef detect_anomalous_params(loss: torch.Tensor, model) -> None:\n    parameters_in_graph = set()\n    visited = set()\n\n    def traverse(grad_fn):\n        if grad_fn is None:\n            return\n        if grad_fn not in visited:\n            visited.add(grad_fn)\n            if hasattr(grad_fn, 'variable'):\n                parameters_in_graph.add(grad_fn.variable)\n            parents = grad_fn.next_functions\n            if parents is not None:\n                for parent in parents:\n                    grad_fn = parent[0]\n                    traverse(grad_fn)\n\n    traverse(loss.grad_fn)\n    from mmengine.logging import MMLogger\n    logger = MMLogger.get_current_instance()\n    for n, p in model.named_parameters():\n        if p not in parameters_in_graph and p.requires_grad:\n            logger.log(\n                level=logging.ERROR,\n                msg=f'{n} with shape {p.size()} is not '\n                f'in the computational graph \\n')\n\n\ndef merge_dict(*args):\n    \"\"\"Merge all dictionaries into one dictionary.\n\n    If pytorch version >= 1.8, ``merge_dict`` will be wrapped\n    by ``torch.fx.wrap``,  which will make ``torch.fx.symbolic_trace`` skip\n    trace ``merge_dict``.\n\n    Note:\n        If a function needs to be traced by ``torch.fx.symbolic_trace``,\n        but inevitably needs to use ``update`` method of ``dict``(``update``\n        is not traceable). It should use ``merge_dict`` to replace\n        ``xxx.update``.\n\n    Args:\n        *args: dictionary needs to be merged.\n\n    Returns:\n        dict: Merged dict from args\n    \"\"\"\n    output = dict()\n    for item in args:\n        assert isinstance(\n            item,\n            dict), (f'all arguments of merge_dict should be a dict, but got '\n                    f'{type(item)}')\n        output.update(item)\n    return output\n\n\n# torch.fx is only available when pytorch version >= 1.8.\n# If the subclass of `BaseModel` has multiple submodules, and each module\n# will return a loss dict during training process, i.e., `TwoStageDetector`\n# in mmdet. It should use `merge_dict` to get the total loss, rather than\n# `loss.update` to keep model traceable.\ntry:\n    import torch.fx\n\n    # make torch.fx skip trace `merge_dict`.\n    merge_dict = torch.fx.wrap(merge_dict)\n\nexcept ImportError:\n    warnings.warn('Cannot import torch.fx, `merge_dict` is a simple function '\n                  'to merge multiple dicts')\n\n\nclass _BatchNormXd(nn.modules.batchnorm._BatchNorm):\n    \"\"\"A general BatchNorm layer without input dimension check.\n\n    Reproduced from @kapily's work:\n    (https://github.com/pytorch/pytorch/issues/41081#issuecomment-783961547)\n    The only difference between BatchNorm1d, BatchNorm2d, BatchNorm3d, etc\n    is `_check_input_dim` that is designed for tensor sanity checks.\n    The check has been bypassed in this class for the convenience of converting\n    SyncBatchNorm.\n    \"\"\"\n\n    def _check_input_dim(self, input: torch.Tensor):\n        return\n\n\ndef revert_sync_batchnorm(module: nn.Module) -> nn.Module:\n    \"\"\"Helper function to convert all `SyncBatchNorm` (SyncBN) and\n    `mmcv.ops.sync_bn.SyncBatchNorm`(MMSyncBN) layers in the model to\n    `BatchNormXd` layers.\n\n    Adapted from @kapily's work:\n    (https://github.com/pytorch/pytorch/issues/41081#issuecomment-783961547)\n\n    Args:\n        module (nn.Module): The module containing `SyncBatchNorm` layers.\n\n    Returns:\n        module_output: The converted module with `BatchNormXd` layers.\n    \"\"\"\n    module_output = module\n    module_checklist = [torch.nn.modules.batchnorm.SyncBatchNorm]\n\n    if mmcv_full_available():\n        from mmcv.ops import SyncBatchNorm\n        module_checklist.append(SyncBatchNorm)\n\n    if isinstance(module, tuple(module_checklist)):\n        module_output = _BatchNormXd(module.num_features, module.eps,\n                                     module.momentum, module.affine,\n                                     module.track_running_stats)\n        if module.affine:\n            # no_grad() may not be needed here but\n            # just to be consistent with `convert_sync_batchnorm()`\n            with torch.no_grad():\n                module_output.weight = module.weight\n                module_output.bias = module.bias\n        module_output.running_mean = module.running_mean\n        module_output.running_var = module.running_var\n        module_output.num_batches_tracked = module.num_batches_tracked\n        module_output.training = module.training\n        # qconfig exists in quantized models\n        if hasattr(module, 'qconfig'):\n            module_output.qconfig = module.qconfig\n    for name, child in module.named_children():\n        # Some custom modules or 3rd party implemented modules may raise an\n        # error when calling `add_module`. Therefore, try to catch the error\n        # and do not raise it. See https://github.com/open-mmlab/mmengine/issues/638 # noqa: E501\n        # for more details.\n        try:\n            module_output.add_module(name, revert_sync_batchnorm(child))\n        except Exception:\n            print_log(\n                F'Failed to convert {child} from SyncBN to BN!',\n                logger='current',\n                level=logging.WARNING)\n    del module\n    return module_output\n\n\ndef convert_sync_batchnorm(module: nn.Module,\n                           implementation='torch') -> nn.Module:\n    \"\"\"Helper function to convert all `BatchNorm` layers in the model to\n    `SyncBatchNorm` (SyncBN) or `mmcv.ops.sync_bn.SyncBatchNorm`(MMSyncBN)\n    layers. Adapted from <https://pytorch.org/docs/stable/generated/torch.nn.Sy\n    ncBatchNorm.html#torch.nn.SyncBatchNorm.convert_sync_batchnorm>_.\n\n    Args:\n        module (nn.Module): The module containing `SyncBatchNorm` layers.\n        implementation (str): The type of `SyncBatchNorm` to convert to.\n\n            - 'torch': convert to `torch.nn.modules.batchnorm.SyncBatchNorm`.\n            - 'mmcv': convert to `mmcv.ops.sync_bn.SyncBatchNorm`.\n\n    Returns:\n        nn.Module: The converted module with `SyncBatchNorm` layers.\n    \"\"\"  # noqa: E501\n    module_output = module\n\n    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):\n        if implementation == 'torch':\n            SyncBatchNorm = torch.nn.modules.batchnorm.SyncBatchNorm\n        elif implementation == 'mmcv':\n            from mmcv.ops import SyncBatchNorm  # type: ignore\n        else:\n            raise ValueError('sync_bn should be \"torch\" or \"mmcv\", but got '\n                             f'{implementation}')\n\n        module_output = SyncBatchNorm(module.num_features, module.eps,\n                                      module.momentum, module.affine,\n                                      module.track_running_stats)\n\n        if module.affine:\n            with torch.no_grad():\n                module_output.weight = module.weight\n                module_output.bias = module.bias\n        module_output.running_mean = module.running_mean\n        module_output.running_var = module.running_var\n        module_output.num_batches_tracked = module.num_batches_tracked\n        if hasattr(module, 'qconfig'):\n            module_output.qconfig = module.qconfig\n    for name, child in module.named_children():\n        module_output.add_module(name,\n                                 convert_sync_batchnorm(child, implementation))\n    del module\n    return module_output\n", "path": "mmengine/model/utils.py"}]}