We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Got unexpected keyword argument 'channel_last' when using apex.parallel.convert_syncbn_model
As https://github.com/NVIDIA/apex/blob/b2f63c48408e1110bb1c7ad7bf6310141da30616/apex/parallel/__init__.py#L41 shows, `apex.parallel.convert_syncbn_model` pass `channel_last` to `SyncBatchNorm`, but `SyncBatchNorm` does not accept. I'm not sure why this function adds a `channel_last` keyword, PyTorch uses channel first tensor, while TensorFlow uses channel last. Simply remove channel_last from the function call fix this issue.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `apex/parallel/sync_batchnorm.py`
Content:
```
1 import torch
2 from torch.nn.modules.batchnorm import _BatchNorm
3 from torch.nn import functional as F
4 
5 from .sync_batchnorm_kernel import SyncBatchnormFunction
6 from apex.parallel import ReduceOp
7 
8 
9 class SyncBatchNorm(_BatchNorm):
10     """
11     synchronized batch normalization module extented from ``torch.nn.BatchNormNd``
12     with the added stats reduction across multiple processes.
13     :class:`apex.parallel.SyncBatchNorm` is designed to work with
14     ``DistributedDataParallel``.
15 
16     When running in training mode, the layer reduces stats across all processes
17     to increase the effective batchsize for normalization layer. This is useful
18     in applications where batch size is small on a given process that would
19     diminish converged accuracy of the model. The model uses collective
20     communication package from ``torch.distributed``.
21 
22     When running in evaluation mode, the layer falls back to
23     ``torch.nn.functional.batch_norm``.
24 
25     Args:
26         num_features: :math:`C` from an expected input of size
27             :math:`(N, C, L)` or :math:`L` from input of size :math:`(N, L)`
28         eps: a value added to the denominator for numerical stability.
29             Default: 1e-5
30         momentum: the value used for the running_mean and running_var
31             computation. Can be set to ``None`` for cumulative moving average
32             (i.e. simple average). Default: 0.1
33         affine: a boolean value that when set to ``True``, this module has
34             learnable affine parameters. Default: ``True``
35         track_running_stats: a boolean value that when set to ``True``, this
36             module tracks the running mean and variance, and when set to ``False``,
37             this module does not track such statistics and always uses batch
38             statistics in both training and eval modes. Default: ``True``
39 
40     Example::
41 
42         >>> sbn = apex.parallel.SyncBatchNorm(100).cuda()
43         >>> inp = torch.randn(10, 100, 14, 14).cuda()
44         >>> out = sbn(inp)
45         >>> inp = torch.randn(3, 100, 20).cuda()
46         >>> out = sbn(inp)
47     """
48 
49     warned = False
50 
51     def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None):
52 
53         if not SyncBatchNorm.warned:
54             print("Warning:  using Python fallback for SyncBatchNorm, possibly because apex was installed without --cuda_ext.  The exception raised when attempting to import the cuda backend was: ", self.syncbn_import_error)
55             SyncBatchNorm.warned = True
56 
57         super(SyncBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)
58         self.process_group = process_group
59 
60     def _specify_process_group(self, process_group):
61         self.process_group = process_group
62 
63     def forward(self, input):
64         torch.cuda.nvtx.range_push("sync_bn_fw_with_mean_var")
65         mean = None
66         var = None
67         if not self.training and self.track_running_stats:
68             # fall back to pytorch implementation for inference
69             torch.cuda.nvtx.range_pop()
70             return F.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps)
71         else:
72             process_group = self.process_group
73             world_size = 0
74             if not self.process_group:
75                 process_group = torch.distributed.group.WORLD
76             world_size = torch.distributed.get_world_size(process_group)
77             self.num_batches_tracked += 1
78             with torch.no_grad():
79                 channel_first_input = input.transpose(0, 1).contiguous()
80                 squashed_input_tensor_view = channel_first_input.view(
81                     channel_first_input.size(0), -1)
82                 # total number of data points for each variance entry. Used to calculate unbiased variance estimate
83                 m = None
84                 local_m = float(squashed_input_tensor_view.size()[1])
85                 local_mean = torch.mean(squashed_input_tensor_view, 1)
86                 local_sqr_mean = torch.pow(
87                     squashed_input_tensor_view, 2).mean(1)
88                 if torch.distributed.is_initialized():
89                     torch.distributed.all_reduce(
90                         local_mean, ReduceOp.SUM, process_group)
91                     mean = local_mean / world_size
92                     torch.distributed.all_reduce(
93                         local_sqr_mean, ReduceOp.SUM, process_group)
94                     sqr_mean = local_sqr_mean / world_size
95                     m = local_m * world_size
96                 else:
97                     m = local_m
98                     mean = local_mean
99                     sqr_mean = local_sqr_mean
100                 # var(x) = E (( x - mean_x ) ** 2)
101                 #        = 1 / N * sum ( x - mean_x ) ** 2
102                 #        = 1 / N * sum (x**2) - mean_x**2
103                 var = sqr_mean - mean.pow(2)
104 
105                 if self.running_mean is not None:
106                     self.running_mean = self.momentum * mean + \
107                         (1 - self.momentum) * self.running_mean
108                 if self.running_var is not None:
109                     # as noted by the paper, we used unbiased variance estimate of the mini-batch
110                     # Var[x] = m / (m-1) * Eb (sample_variance)
111                     self.running_var = m / \
112                         (m-1) * self.momentum * var + \
113                         (1 - self.momentum) * self.running_var
114             torch.cuda.nvtx.range_pop()
115             return SyncBatchnormFunction.apply(input, self.weight, self.bias, mean, var, self.eps, process_group, world_size)
116
```

Path: `apex/parallel/optimized_sync_batchnorm.py`
Content:
```
1 import torch
2 from torch.nn.modules.batchnorm import _BatchNorm
3 from torch.nn import functional as F
4 
5 import syncbn
6 from .optimized_sync_batchnorm_kernel import SyncBatchnormFunction
7 
8 
9 class SyncBatchNorm(_BatchNorm):
10     """
11     synchronized batch normalization module extented from `torch.nn.BatchNormNd`
12     with the added stats reduction across multiple processes.
13     :class:`apex.parallel.SyncBatchNorm` is designed to work with
14     `DistributedDataParallel`.
15 
16     When running in training mode, the layer reduces stats across all processes
17     to increase the effective batchsize for normalization layer. This is useful
18     in applications where batch size is small on a given process that would
19     diminish converged accuracy of the model. The model uses collective
20     communication package from `torch.distributed`.
21 
22     When running in evaluation mode, the layer falls back to
23     `torch.nn.functional.batch_norm`
24 
25     Args:
26         num_features: :math:`C` from an expected input of size
27             :math:`(N, C, L)` or :math:`L` from input of size :math:`(N, L)`
28         eps: a value added to the denominator for numerical stability.
29             Default: 1e-5
30         momentum: the value used for the running_mean and running_var
31             computation. Can be set to ``None`` for cumulative moving average
32             (i.e. simple average). Default: 0.1
33         affine: a boolean value that when set to ``True``, this module has
34             learnable affine parameters. Default: ``True``
35         track_running_stats: a boolean value that when set to ``True``, this
36             module tracks the running mean and variance, and when set to ``False``,
37             this module does not track such statistics and always uses batch
38             statistics in both training and eval modes. Default: ``True``
39         process_group: pass in a process group within which the stats of the
40             mini-batch is being synchronized. ``None`` for using default process
41             group
42         channel_last: a boolean value that when set to ``True``, this module
43             take the last dimension of the input tensor to be the channel
44             dimension. Default: False
45 
46     Examples::
47         >>> # channel first tensor
48         >>> sbn = apex.parallel.SyncBatchNorm(100).cuda()
49         >>> inp = torch.randn(10, 100, 14, 14).cuda()
50         >>> out = sbn(inp)
51         >>> inp = torch.randn(3, 100, 20).cuda()
52         >>> out = sbn(inp)
53         >>> # channel last tensor
54         >>> sbn = apex.parallel.SyncBatchNorm(100, channel_last=True).cuda()
55         >>> inp = torch.randn(10, 14, 14, 100).cuda()
56     """
57 
58     def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last = False):
59         super(SyncBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)
60         self.process_group = process_group
61         self.channel_last = channel_last
62 
63     def _specify_process_group(self, process_group):
64         self.process_group = process_group
65 
66     def _specify_channel_last(self, channel_last):
67         self.channel_last = channel_last
68 
69     def forward(self, input):
70         if not self.training and self.track_running_stats and not self.channel_last:
71             # fall back to pytorch implementation for inference
72             return F.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps)
73         else:
74             exponential_average_factor = 0.0
75             if self.training and self.track_running_stats:
76                 self.num_batches_tracked += 1
77                 if self.momentum is None:
78                     exponential_average_factor = 1.0 / float(self.num_batches_tracked)
79                 else:
80                     exponential_average_factor = self.momentum
81             return SyncBatchnormFunction.apply(input, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.training or not self.track_running_stats, exponential_average_factor, self.process_group, self.channel_last)
82
```
--- 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/apex/parallel/optimized_sync_batchnorm.py b/apex/parallel/optimized_sync_batchnorm.py
--- a/apex/parallel/optimized_sync_batchnorm.py
+++ b/apex/parallel/optimized_sync_batchnorm.py
@@ -55,7 +55,7 @@
         >>> inp = torch.randn(10, 14, 14, 100).cuda()
     """
 
-    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last = False):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last=False):
         super(SyncBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)
         self.process_group = process_group
         self.channel_last = channel_last
diff --git a/apex/parallel/sync_batchnorm.py b/apex/parallel/sync_batchnorm.py
--- a/apex/parallel/sync_batchnorm.py
+++ b/apex/parallel/sync_batchnorm.py
@@ -48,7 +48,9 @@
 
     warned = False
 
-    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last=False):
+        if channel_last == True:
+            raise AttributeError("channel_last is not supported by primitive SyncBatchNorm implementation. Try install apex with `--cuda_ext` if channel_last is desired.")
 
         if not SyncBatchNorm.warned:
             print("Warning:  using Python fallback for SyncBatchNorm, possibly because apex was installed without --cuda_ext.  The exception raised when attempting to import the cuda backend was: ", self.syncbn_import_error)

{"golden_diff": "diff --git a/apex/parallel/optimized_sync_batchnorm.py b/apex/parallel/optimized_sync_batchnorm.py\n--- a/apex/parallel/optimized_sync_batchnorm.py\n+++ b/apex/parallel/optimized_sync_batchnorm.py\n@@ -55,7 +55,7 @@\n         >>> inp = torch.randn(10, 14, 14, 100).cuda()\n     \"\"\"\n \n-    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last = False):\n+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last=False):\n         super(SyncBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)\n         self.process_group = process_group\n         self.channel_last = channel_last\ndiff --git a/apex/parallel/sync_batchnorm.py b/apex/parallel/sync_batchnorm.py\n--- a/apex/parallel/sync_batchnorm.py\n+++ b/apex/parallel/sync_batchnorm.py\n@@ -48,7 +48,9 @@\n \n     warned = False\n \n-    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None):\n+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last=False):\n+        if channel_last == True:\n+            raise AttributeError(\"channel_last is not supported by primitive SyncBatchNorm implementation. Try install apex with `--cuda_ext` if channel_last is desired.\")\n \n         if not SyncBatchNorm.warned:\n             print(\"Warning:  using Python fallback for SyncBatchNorm, possibly because apex was installed without --cuda_ext.  The exception raised when attempting to import the cuda backend was: \", self.syncbn_import_error)\n", "issue": "Got unexpected keyword argument 'channel_last' when using apex.parallel.convert_syncbn_model\nAs https://github.com/NVIDIA/apex/blob/b2f63c48408e1110bb1c7ad7bf6310141da30616/apex/parallel/__init__.py#L41 shows, `apex.parallel.convert_syncbn_model` pass `channel_last` to `SyncBatchNorm`, but `SyncBatchNorm` does not accept. I'm not sure why this function adds a `channel_last` keyword, PyTorch uses channel first tensor, while TensorFlow uses channel last. Simply remove channel_last from the function call fix this issue.\n", "before_files": [{"content": "import torch\nfrom torch.nn.modules.batchnorm import _BatchNorm\nfrom torch.nn import functional as F\n\nfrom .sync_batchnorm_kernel import SyncBatchnormFunction\nfrom apex.parallel import ReduceOp\n\n\nclass SyncBatchNorm(_BatchNorm):\n    \"\"\"\n    synchronized batch normalization module extented from ``torch.nn.BatchNormNd``\n    with the added stats reduction across multiple processes.\n    :class:`apex.parallel.SyncBatchNorm` is designed to work with\n    ``DistributedDataParallel``.\n\n    When running in training mode, the layer reduces stats across all processes\n    to increase the effective batchsize for normalization layer. This is useful\n    in applications where batch size is small on a given process that would\n    diminish converged accuracy of the model. The model uses collective\n    communication package from ``torch.distributed``.\n\n    When running in evaluation mode, the layer falls back to\n    ``torch.nn.functional.batch_norm``.\n\n    Args:\n        num_features: :math:`C` from an expected input of size\n            :math:`(N, C, L)` or :math:`L` from input of size :math:`(N, L)`\n        eps: a value added to the denominator for numerical stability.\n            Default: 1e-5\n        momentum: the value used for the running_mean and running_var\n            computation. Can be set to ``None`` for cumulative moving average\n            (i.e. simple average). Default: 0.1\n        affine: a boolean value that when set to ``True``, this module has\n            learnable affine parameters. Default: ``True``\n        track_running_stats: a boolean value that when set to ``True``, this\n            module tracks the running mean and variance, and when set to ``False``,\n            this module does not track such statistics and always uses batch\n            statistics in both training and eval modes. Default: ``True``\n\n    Example::\n\n        >>> sbn = apex.parallel.SyncBatchNorm(100).cuda()\n        >>> inp = torch.randn(10, 100, 14, 14).cuda()\n        >>> out = sbn(inp)\n        >>> inp = torch.randn(3, 100, 20).cuda()\n        >>> out = sbn(inp)\n    \"\"\"\n\n    warned = False\n\n    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None):\n\n        if not SyncBatchNorm.warned:\n            print(\"Warning:  using Python fallback for SyncBatchNorm, possibly because apex was installed without --cuda_ext.  The exception raised when attempting to import the cuda backend was: \", self.syncbn_import_error)\n            SyncBatchNorm.warned = True\n\n        super(SyncBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)\n        self.process_group = process_group\n\n    def _specify_process_group(self, process_group):\n        self.process_group = process_group\n\n    def forward(self, input):\n        torch.cuda.nvtx.range_push(\"sync_bn_fw_with_mean_var\")\n        mean = None\n        var = None\n        if not self.training and self.track_running_stats:\n            # fall back to pytorch implementation for inference\n            torch.cuda.nvtx.range_pop()\n            return F.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps)\n        else:\n            process_group = self.process_group\n            world_size = 0\n            if not self.process_group:\n                process_group = torch.distributed.group.WORLD\n            world_size = torch.distributed.get_world_size(process_group)\n            self.num_batches_tracked += 1\n            with torch.no_grad():\n                channel_first_input = input.transpose(0, 1).contiguous()\n                squashed_input_tensor_view = channel_first_input.view(\n                    channel_first_input.size(0), -1)\n                # total number of data points for each variance entry. Used to calculate unbiased variance estimate\n                m = None\n                local_m = float(squashed_input_tensor_view.size()[1])\n                local_mean = torch.mean(squashed_input_tensor_view, 1)\n                local_sqr_mean = torch.pow(\n                    squashed_input_tensor_view, 2).mean(1)\n                if torch.distributed.is_initialized():\n                    torch.distributed.all_reduce(\n                        local_mean, ReduceOp.SUM, process_group)\n                    mean = local_mean / world_size\n                    torch.distributed.all_reduce(\n                        local_sqr_mean, ReduceOp.SUM, process_group)\n                    sqr_mean = local_sqr_mean / world_size\n                    m = local_m * world_size\n                else:\n                    m = local_m\n                    mean = local_mean\n                    sqr_mean = local_sqr_mean\n                # var(x) = E (( x - mean_x ) ** 2)\n                #        = 1 / N * sum ( x - mean_x ) ** 2\n                #        = 1 / N * sum (x**2) - mean_x**2\n                var = sqr_mean - mean.pow(2)\n\n                if self.running_mean is not None:\n                    self.running_mean = self.momentum * mean + \\\n                        (1 - self.momentum) * self.running_mean\n                if self.running_var is not None:\n                    # as noted by the paper, we used unbiased variance estimate of the mini-batch\n                    # Var[x] = m / (m-1) * Eb (sample_variance)\n                    self.running_var = m / \\\n                        (m-1) * self.momentum * var + \\\n                        (1 - self.momentum) * self.running_var\n            torch.cuda.nvtx.range_pop()\n            return SyncBatchnormFunction.apply(input, self.weight, self.bias, mean, var, self.eps, process_group, world_size)\n", "path": "apex/parallel/sync_batchnorm.py"}, {"content": "import torch\nfrom torch.nn.modules.batchnorm import _BatchNorm\nfrom torch.nn import functional as F\n\nimport syncbn\nfrom .optimized_sync_batchnorm_kernel import SyncBatchnormFunction\n\n\nclass SyncBatchNorm(_BatchNorm):\n    \"\"\"\n    synchronized batch normalization module extented from `torch.nn.BatchNormNd`\n    with the added stats reduction across multiple processes.\n    :class:`apex.parallel.SyncBatchNorm` is designed to work with\n    `DistributedDataParallel`.\n\n    When running in training mode, the layer reduces stats across all processes\n    to increase the effective batchsize for normalization layer. This is useful\n    in applications where batch size is small on a given process that would\n    diminish converged accuracy of the model. The model uses collective\n    communication package from `torch.distributed`.\n\n    When running in evaluation mode, the layer falls back to\n    `torch.nn.functional.batch_norm`\n\n    Args:\n        num_features: :math:`C` from an expected input of size\n            :math:`(N, C, L)` or :math:`L` from input of size :math:`(N, L)`\n        eps: a value added to the denominator for numerical stability.\n            Default: 1e-5\n        momentum: the value used for the running_mean and running_var\n            computation. Can be set to ``None`` for cumulative moving average\n            (i.e. simple average). Default: 0.1\n        affine: a boolean value that when set to ``True``, this module has\n            learnable affine parameters. Default: ``True``\n        track_running_stats: a boolean value that when set to ``True``, this\n            module tracks the running mean and variance, and when set to ``False``,\n            this module does not track such statistics and always uses batch\n            statistics in both training and eval modes. Default: ``True``\n        process_group: pass in a process group within which the stats of the\n            mini-batch is being synchronized. ``None`` for using default process\n            group\n        channel_last: a boolean value that when set to ``True``, this module\n            take the last dimension of the input tensor to be the channel\n            dimension. Default: False\n\n    Examples::\n        >>> # channel first tensor\n        >>> sbn = apex.parallel.SyncBatchNorm(100).cuda()\n        >>> inp = torch.randn(10, 100, 14, 14).cuda()\n        >>> out = sbn(inp)\n        >>> inp = torch.randn(3, 100, 20).cuda()\n        >>> out = sbn(inp)\n        >>> # channel last tensor\n        >>> sbn = apex.parallel.SyncBatchNorm(100, channel_last=True).cuda()\n        >>> inp = torch.randn(10, 14, 14, 100).cuda()\n    \"\"\"\n\n    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last = False):\n        super(SyncBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)\n        self.process_group = process_group\n        self.channel_last = channel_last\n\n    def _specify_process_group(self, process_group):\n        self.process_group = process_group\n\n    def _specify_channel_last(self, channel_last):\n        self.channel_last = channel_last\n\n    def forward(self, input):\n        if not self.training and self.track_running_stats and not self.channel_last:\n            # fall back to pytorch implementation for inference\n            return F.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps)\n        else:\n            exponential_average_factor = 0.0\n            if self.training and self.track_running_stats:\n                self.num_batches_tracked += 1\n                if self.momentum is None:\n                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)\n                else:\n                    exponential_average_factor = self.momentum\n            return SyncBatchnormFunction.apply(input, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.training or not self.track_running_stats, exponential_average_factor, self.process_group, self.channel_last)\n", "path": "apex/parallel/optimized_sync_batchnorm.py"}], "after_files": [{"content": "import torch\nfrom torch.nn.modules.batchnorm import _BatchNorm\nfrom torch.nn import functional as F\n\nfrom .sync_batchnorm_kernel import SyncBatchnormFunction\nfrom apex.parallel import ReduceOp\n\n\nclass SyncBatchNorm(_BatchNorm):\n    \"\"\"\n    synchronized batch normalization module extented from ``torch.nn.BatchNormNd``\n    with the added stats reduction across multiple processes.\n    :class:`apex.parallel.SyncBatchNorm` is designed to work with\n    ``DistributedDataParallel``.\n\n    When running in training mode, the layer reduces stats across all processes\n    to increase the effective batchsize for normalization layer. This is useful\n    in applications where batch size is small on a given process that would\n    diminish converged accuracy of the model. The model uses collective\n    communication package from ``torch.distributed``.\n\n    When running in evaluation mode, the layer falls back to\n    ``torch.nn.functional.batch_norm``.\n\n    Args:\n        num_features: :math:`C` from an expected input of size\n            :math:`(N, C, L)` or :math:`L` from input of size :math:`(N, L)`\n        eps: a value added to the denominator for numerical stability.\n            Default: 1e-5\n        momentum: the value used for the running_mean and running_var\n            computation. Can be set to ``None`` for cumulative moving average\n            (i.e. simple average). Default: 0.1\n        affine: a boolean value that when set to ``True``, this module has\n            learnable affine parameters. Default: ``True``\n        track_running_stats: a boolean value that when set to ``True``, this\n            module tracks the running mean and variance, and when set to ``False``,\n            this module does not track such statistics and always uses batch\n            statistics in both training and eval modes. Default: ``True``\n\n    Example::\n\n        >>> sbn = apex.parallel.SyncBatchNorm(100).cuda()\n        >>> inp = torch.randn(10, 100, 14, 14).cuda()\n        >>> out = sbn(inp)\n        >>> inp = torch.randn(3, 100, 20).cuda()\n        >>> out = sbn(inp)\n    \"\"\"\n\n    warned = False\n\n    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last=False):\n        if channel_last == True:\n            raise AttributeError(\"channel_last is not supported by primitive SyncBatchNorm implementation. Try install apex with `--cuda_ext` if channel_last is desired.\")\n\n        if not SyncBatchNorm.warned:\n            print(\"Warning:  using Python fallback for SyncBatchNorm, possibly because apex was installed without --cuda_ext.  The exception raised when attempting to import the cuda backend was: \", self.syncbn_import_error)\n            SyncBatchNorm.warned = True\n\n        super(SyncBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)\n        self.process_group = process_group\n\n    def _specify_process_group(self, process_group):\n        self.process_group = process_group\n\n    def forward(self, input):\n        torch.cuda.nvtx.range_push(\"sync_bn_fw_with_mean_var\")\n        mean = None\n        var = None\n        if not self.training and self.track_running_stats:\n            # fall back to pytorch implementation for inference\n            torch.cuda.nvtx.range_pop()\n            return F.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps)\n        else:\n            process_group = self.process_group\n            world_size = 0\n            if not self.process_group:\n                process_group = torch.distributed.group.WORLD\n            world_size = torch.distributed.get_world_size(process_group)\n            self.num_batches_tracked += 1\n            with torch.no_grad():\n                channel_first_input = input.transpose(0, 1).contiguous()\n                squashed_input_tensor_view = channel_first_input.view(\n                    channel_first_input.size(0), -1)\n                # total number of data points for each variance entry. Used to calculate unbiased variance estimate\n                m = None\n                local_m = float(squashed_input_tensor_view.size()[1])\n                local_mean = torch.mean(squashed_input_tensor_view, 1)\n                local_sqr_mean = torch.pow(\n                    squashed_input_tensor_view, 2).mean(1)\n                if torch.distributed.is_initialized():\n                    torch.distributed.all_reduce(\n                        local_mean, ReduceOp.SUM, process_group)\n                    mean = local_mean / world_size\n                    torch.distributed.all_reduce(\n                        local_sqr_mean, ReduceOp.SUM, process_group)\n                    sqr_mean = local_sqr_mean / world_size\n                    m = local_m * world_size\n                else:\n                    m = local_m\n                    mean = local_mean\n                    sqr_mean = local_sqr_mean\n                # var(x) = E (( x - mean_x ) ** 2)\n                #        = 1 / N * sum ( x - mean_x ) ** 2\n                #        = 1 / N * sum (x**2) - mean_x**2\n                var = sqr_mean - mean.pow(2)\n\n                if self.running_mean is not None:\n                    self.running_mean = self.momentum * mean + \\\n                        (1 - self.momentum) * self.running_mean\n                if self.running_var is not None:\n                    # as noted by the paper, we used unbiased variance estimate of the mini-batch\n                    # Var[x] = m / (m-1) * Eb (sample_variance)\n                    self.running_var = m / \\\n                        (m-1) * self.momentum * var + \\\n                        (1 - self.momentum) * self.running_var\n            torch.cuda.nvtx.range_pop()\n            return SyncBatchnormFunction.apply(input, self.weight, self.bias, mean, var, self.eps, process_group, world_size)\n", "path": "apex/parallel/sync_batchnorm.py"}, {"content": "import torch\nfrom torch.nn.modules.batchnorm import _BatchNorm\nfrom torch.nn import functional as F\n\nimport syncbn\nfrom .optimized_sync_batchnorm_kernel import SyncBatchnormFunction\n\n\nclass SyncBatchNorm(_BatchNorm):\n    \"\"\"\n    synchronized batch normalization module extented from `torch.nn.BatchNormNd`\n    with the added stats reduction across multiple processes.\n    :class:`apex.parallel.SyncBatchNorm` is designed to work with\n    `DistributedDataParallel`.\n\n    When running in training mode, the layer reduces stats across all processes\n    to increase the effective batchsize for normalization layer. This is useful\n    in applications where batch size is small on a given process that would\n    diminish converged accuracy of the model. The model uses collective\n    communication package from `torch.distributed`.\n\n    When running in evaluation mode, the layer falls back to\n    `torch.nn.functional.batch_norm`\n\n    Args:\n        num_features: :math:`C` from an expected input of size\n            :math:`(N, C, L)` or :math:`L` from input of size :math:`(N, L)`\n        eps: a value added to the denominator for numerical stability.\n            Default: 1e-5\n        momentum: the value used for the running_mean and running_var\n            computation. Can be set to ``None`` for cumulative moving average\n            (i.e. simple average). Default: 0.1\n        affine: a boolean value that when set to ``True``, this module has\n            learnable affine parameters. Default: ``True``\n        track_running_stats: a boolean value that when set to ``True``, this\n            module tracks the running mean and variance, and when set to ``False``,\n            this module does not track such statistics and always uses batch\n            statistics in both training and eval modes. Default: ``True``\n        process_group: pass in a process group within which the stats of the\n            mini-batch is being synchronized. ``None`` for using default process\n            group\n        channel_last: a boolean value that when set to ``True``, this module\n            take the last dimension of the input tensor to be the channel\n            dimension. Default: False\n\n    Examples::\n        >>> # channel first tensor\n        >>> sbn = apex.parallel.SyncBatchNorm(100).cuda()\n        >>> inp = torch.randn(10, 100, 14, 14).cuda()\n        >>> out = sbn(inp)\n        >>> inp = torch.randn(3, 100, 20).cuda()\n        >>> out = sbn(inp)\n        >>> # channel last tensor\n        >>> sbn = apex.parallel.SyncBatchNorm(100, channel_last=True).cuda()\n        >>> inp = torch.randn(10, 14, 14, 100).cuda()\n    \"\"\"\n\n    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last=False):\n        super(SyncBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)\n        self.process_group = process_group\n        self.channel_last = channel_last\n\n    def _specify_process_group(self, process_group):\n        self.process_group = process_group\n\n    def _specify_channel_last(self, channel_last):\n        self.channel_last = channel_last\n\n    def forward(self, input):\n        if not self.training and self.track_running_stats and not self.channel_last:\n            # fall back to pytorch implementation for inference\n            return F.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps)\n        else:\n            exponential_average_factor = 0.0\n            if self.training and self.track_running_stats:\n                self.num_batches_tracked += 1\n                if self.momentum is None:\n                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)\n                else:\n                    exponential_average_factor = self.momentum\n            return SyncBatchnormFunction.apply(input, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.training or not self.track_running_stats, exponential_average_factor, self.process_group, self.channel_last)\n", "path": "apex/parallel/optimized_sync_batchnorm.py"}]}