Repository: NVIDIA/apex

Issue:
SyncBatchNorm doesn't support 2 dimensions input?
Hi,
I'm facing the issue that the program crash when the input for SyncBatchNorm is two dimensions. Here's the code:
```python
import torch
import apex

model = apex.parallel.SyncBatchNorm(4).cuda()
data = torch.rand((8,4)).cuda()
output = model(data)
```
When running the code, error raised like this:
```
Traceback (most recent call last):
  File "syncbn_test.by", line 7, in <module>
    output = model(data)
  File "/usr/local/lib/python3.5/dist-packages/torch/nn/modules/module.py", line 489, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.5/dist-packages/apex/parallel/optimized_sync_batchnorm.py", line 81, in forward
    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)
  File "/usr/local/lib/python3.5/dist-packages/apex/parallel/optimized_sync_batchnorm_kernel.py", line 27, in forward
    mean, var_biased = syncbn.welford_mean_var(input)
RuntimeError: Dimension out of range (expected to be in range of [-2, 1], but got 2) (maybe_wrap_dim at /pytorch/aten/src/ATen/core/WrapDimMinimal.h:18)
```
And everthing runs ok when `data` a 4 dims tensor. 

Here is my environment:
```
Ubuntu 16.04
Python 3.5.2
Pytorch 1.0.1, installed with "pip install torch"
apex is installed with command:
  pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .
cuda 10.0
nvidia driver 410.72
```


Discussion:
@jjsjann123 have you observed this?
yeah, this is an overlook for my side.
I mistakenly thought this is not supported by torch.nn.BatchNormXd, but it is.

I'll patch it soon

Thanks very much. Hope for reply once fixed!
I'll update on this thread once it's fixed, should be sometime this week (Currently occupied by some other work)
I'm facing the same issue, hoping you guys can fix this soon..... many thanks!
The support should be really easy as we already have channel_last kernels there, that handles the same data stride as with 2 dimensional tensor.
I'll update later today or early tomorrow as soon as I finish my work at hand
@flymark2010 Sorry for the long wait for such a simple fix (priority issue on my side :/ )
PR with the fix has been issued. Feel free to cherry pick, I'll push for the merge as well.
Should be fixed via https://github.com/NVIDIA/apex/commit/0a991543846966d5f586540dc2441e512139e9fc.
I am using latest master and got below issue:
![image](https://user-images.githubusercontent.com/4326958/54980485-d1a48700-4fe0-11e9-8cbe-ed9aa5d5dee4.png)
PS:
self.bn5 = SyncBatchNorm(512)

Thanks a lot for reporting this @hiyijian. I should have done a better job unit testing this.

Any chance I can have a repro? If it's something not related to 2-dimension tensor, maybe we want to have a new issue and tag me on that one :)
Hello, I'm confused with the same problem.
The apex was installed on 16 Sep.

import torch
import apex
model = apex.parallel.SyncBatchNorm(4).cuda()
data = torch.rand((8,4)).cuda()
ouput = model(data)

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/gpfs01/user_home/zhuying/.conda/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py", line 547, in __call__
    result = self.forward(*input, **kwargs)
  File "/gpfs01/user_home/zhuying/.conda/envs/py36/lib/python3.6/site-packages/apex/parallel/optimized_sync_batchnorm.py", line 85, in forward
    return SyncBatchnormFunction.apply(input, z, 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, self.fuse_relu)
  File "/gpfs01/user_home/zhuying/.conda/envs/py36/lib/python3.6/site-packages/apex/parallel/optimized_sync_batchnorm_kernel.py", line 27, in forward
    mean, var_biased = syncbn.welford_mean_var(input)
RuntimeError: Dimension out of range (expected to be in range of [-2, 1], but got 2) (maybe_wrap_dim at /tmp/pip-req-build-p5q91txh/c10/core/WrapDimMinimal.h:20)

> Hello, I'm confused with the same problem.
> The apex was installed on 16 Sep.
> 
> import torch
> import apex
> model = apex.parallel.SyncBatchNorm(4).cuda()
> data = torch.rand((8,4)).cuda()
> ouput = model(data)
> 
> Traceback (most recent call last):
> File "", line 1, in
> File "/gpfs01/user_home/zhuying/.conda/envs/py36/lib/python3.6/site-packages/torch/nn/modules/module.py", line 547, in **call**
> result = self.forward(*input, **kwargs)
> File "/gpfs01/user_home/zhuying/.conda/envs/py36/lib/python3.6/site-packages/apex/parallel/optimized_sync_batchnorm.py", line 85, in forward
> return SyncBatchnormFunction.apply(input, z, 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, self.fuse_relu)
> File "/gpfs01/user_home/zhuying/.conda/envs/py36/lib/python3.6/site-packages/apex/parallel/optimized_sync_batchnorm_kernel.py", line 27, in forward
> mean, var_biased = syncbn.welford_mean_var(input)
> RuntimeError: Dimension out of range (expected to be in range of [-2, 1], but got 2) (maybe_wrap_dim at /tmp/pip-req-build-p5q91txh/c10/core/WrapDimMinimal.h:20)

same error
@zhuyingSeu a workaround is using python-fallback version
Same error, and when try to use python-fallback, get another error.
Model is fine when not use syncbn.
```
Traceback (most recent call last):
  File "dist_train.py", line 74, in <module>
    main()
  File "dist_train.py", line 70, in main
    num_query
  File "/root/gg/tyc/processor/processor.py", line 82, in do_train
    scores, feats = model(img, target)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 539, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/apex/parallel/distributed.py", line 560, in forward
    result = self.module(*inputs, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 539, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/apex/amp/_initialize.py", line 197, in new_fwd
    **applier(kwargs, input_caster))
  File "/root/gg/tyc/model/make_model.py", line 230, in forward
    out = self.model.features(x)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 539, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/container.py", line 92, in forward
    input = module(input)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 539, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torchvision/models/densenet.py", line 74, in forward
    new_features = layer(*features)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 539, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torchvision/models/densenet.py", line 50, in forward
    bottleneck_output = bn_function(*prev_features)
  File "/usr/local/lib/python3.6/dist-packages/torchvision/models/densenet.py", line 23, in bn_function
    bottleneck_output = conv(relu(norm(concated_features)))
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 539, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/conv.py", line 343, in forward
    return self.conv2d_forward(input, self.weight)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/conv.py", line 340, in conv2d_forward
    self.padding, self.dilation, self.groups)
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.cuda.HalfTensor) should be the same
```
Oops, somehow missed this thread, let me take a look at this.
It was broken by me earlier at #275 https://github.com/NVIDIA/apex/pull/275/files#diff-f22c2ebe466a49fc0a46a5e7fabb2004R74-R85 :man_facepalming: 

Will push a fix shortly

Resolving patch:
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
@@ -71,7 +71,7 @@ def forward(self, input, z = None):
         # if input.dim() == 2, we switch to channel_last for efficient memory accessing
         channel_last = self.channel_last if input.dim() != 2 else True
 
-        if not self.training and self.track_running_stats and not self.channel_last and not self.fuse_relu and z == None:
+        if not self.training and self.track_running_stats and not channel_last and not self.fuse_relu and z == None:
             # fall back to pytorch implementation for inference
             return F.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps)
         else:
@@ -82,4 +82,4 @@ def forward(self, input, z = None):
                     exponential_average_factor = 1.0 / float(self.num_batches_tracked)
                 else:
                     exponential_average_factor = self.momentum
-            return SyncBatchnormFunction.apply(input, z, 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, self.fuse_relu)
+            return SyncBatchnormFunction.apply(input, z, 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, channel_last, self.fuse_relu)


Test patch:
diff --git a/tests/distributed/synced_batchnorm/test_batchnorm1d.py b/tests/distributed/synced_batchnorm/test_batchnorm1d.py
new file mode 100644
--- /dev/null
+++ b/tests/distributed/synced_batchnorm/test_batchnorm1d.py
@@ -0,0 +1,18 @@
+import torch
+import apex
+
+model = apex.parallel.SyncBatchNorm(4).cuda()
+model.weight.data.uniform_()
+model.bias.data.uniform_()
+data = torch.rand((8,4)).cuda()
+
+model_ref = torch.nn.BatchNorm1d(4).cuda()
+model_ref.load_state_dict(model.state_dict())
+data_ref = data.clone()
+
+output = model(data)
+output_ref = model_ref(data_ref)
+
+assert(output.allclose(output_ref))
+assert(model.running_mean.allclose(model_ref.running_mean))
+assert(model.running_var.allclose(model_ref.running_var))
diff --git a/tests/distributed/synced_batchnorm/unit_test.sh b/tests/distributed/synced_batchnorm/unit_test.sh
--- a/tests/distributed/synced_batchnorm/unit_test.sh
+++ b/tests/distributed/synced_batchnorm/unit_test.sh
@@ -1,5 +1,6 @@
 python python_single_gpu_unit_test.py
 python single_gpu_unit_test.py
+python test_batchnorm1d.py
 python -m torch.distributed.launch --nproc_per_node=2 two_gpu_unit_test.py
 python -m torch.distributed.launch --nproc_per_node=2 two_gpu_unit_test.py --fp16
 #beware, you need a system with at least 4 gpus to test group_size<world_size