SyncBatchNorm.convert_sync_batchnorm() causes ValueError: expected at least 3D input (got 2D input)

username: farleylai
cooked: <p>Not sure if something is missing but isn’t <code>SyncBatchnorm.convert_sync_batchnorm()</code> supposed to convert the module transparently?<br>
However, the following code segment produce <code>ValueError: expected at least 3D input (got 2D input) </code>.<br>
Without the conversion, the forward goes as expected.<br>
Any ideas?</p>
<pre><code class="lang-auto">import os
import torch
from torch import nn

module = torch.nn.Sequential(
           torch.nn.Linear(20, 100),
           torch.nn.BatchNorm1d(100)
         ).cuda()

# creating process group (optional)
# process_ids is a list of int identifying rank ids.
os.environ['RANK'] = '0' 
os.environ['WORLD_SIZE'] = '1' 
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '25791'

process_group = torch.distributed.init_process_group(backend='nccl')
module = nn.SyncBatchNorm.convert_sync_batchnorm(module, process_group)

input = torch.randn(2, 20).cuda()
output = module(input)
print(output.shape)
</code></pre>
<p>The output:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "syncBN.py", line 21, in &lt;module&gt;
    output = module(input)
  File "/home/ml/farleylai/miniconda3/envs/sinet37/lib/python3.7/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/ml/farleylai/miniconda3/envs/sinet37/lib/python3.7/site-packages/torch/nn/modules/container.py", line 92, in forward
    input = module(input)
  File "/home/ml/farleylai/miniconda3/envs/sinet37/lib/python3.7/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/ml/farleylai/miniconda3/envs/sinet37/lib/python3.7/site-packages/torch/nn/modules/batchnorm.py", line 429, in forward
    self._check_input_dim(input)
  File "/home/ml/farleylai/miniconda3/envs/sinet37/lib/python3.7/site-packages/torch/nn/modules/batchnorm.py", line 417, in _check_input_dim
    .format(input.dim()))
ValueError: expected at least 3D input (got 2D input)
</code></pre>
<p>Expected output as w/o conversion:</p>
<pre><code class="lang-auto">torch.Size([2, 100])
</code></pre>
<p>Ubuntu 16.04 with PyTorch 1.3 installed through conda.</p>

username: albanD
cooked: <p>Hi,</p>
<p>The original modules like <code>BatchNorm1d</code> or <code>BatchNorm2d</code> support not having a batch size, so they handle respectively 2d/3d inputs and 3d/4d inputs.<br>
The sync batchnorm has no specialized functions and works for all. But to know which version to use, it must use the number of dimensions of the input (otherwise as you see above, 3d input could be either a batched 1d or an unbatched 2d). And so it only allows having a batch dimension.</p>

username: farleylai
cooked: <p><a class="mention" href="/u/alband">@albanD</a><br>
I looked into the code and found this restriction is imposed by <code>SyncBatchNorm</code>:</p>
<pre><code class="lang-auto">def _check_input_dim(self, input):
        if input.dim() &lt;= 2:
            raise ValueError('expected at least 3D input (got {}D input)'
                             .format(input.dim()))
</code></pre>
<p>This is completely different from the original <code>BatchNorm1d</code> to be wrapped:</p>
<pre><code class="lang-auto">    def _check_input_dim(self, input):
        if input.dim() != 2 and input.dim() != 3:
            raise ValueError('expected 2D or 3D input (got {}D input)'
                             .format(input.dim()))
</code></pre>
<p>I got confused with the code segment that is actually from your API document of <code>SyncBatchNorm</code> using <code>BatchNorm1d</code> for <code>convert_sync_batchnorm</code>.<br>
Why doesn’t <code>SyncBatchNorm</code> explicitly check whether the module to wrap is <code>BatchNorm1d</code> or <code>BatchNorm2d</code> instead of the general <code>_BatchNorm</code> in <code>convert_sync_batchnorm</code>?<br>
If this is not going to work, what is the right way to use <code>convert_sync_batchnorm</code> for those models with <code>BatchNorm1d</code>?</p>

username: albanD
cooked: <blockquote>
<p>If this is not going to work, what is the right way to use  <code>convert_sync_batchnorm</code>  for those models with  <code>BatchNorm1d</code> ?</p>
</blockquote>
<p>I think the fix here is to ensure you always have a batch dimension. Potentially adding an <code>.unsqueeze(0)</code> to your input.</p>
<blockquote>
<p>Then why doesn’t  <code>SyncBatchNorm</code>  explicitly check whether the module to wrap is  <code>BatchNorm1d</code>  or  <code>BatchNorm2d</code>  instead of the general  <code>_BatchNorm</code>  in  <code>convert_sync_batchnorm</code> ?</p>
</blockquote>
<p>This would be a nice addition, we would be happy to merge a PR that adds this feature!</p>

username: farleylai
cooked: <blockquote>
<p>I think the fix here is to ensure you always have a batch dimension. Potentially adding an  <code>.unsqueeze(0)</code>  to your input.</p>
</blockquote>
<p>The example input (2, 20) already contains a batch dim, indicating a batch of two 1D examples.<br>
If we fake the input with <code>unsqueeze(0)</code>, how could it work when there are other modules before <code>BatchNorm1d</code> in the model that may assume the 0th dim must be the batch dim?<br>
After all, the layer(s) before <code>BatchNorm1d</code> can be anything else in general, right?</p>
<p>BTW, I tried to make it of size (1, 2, 20) but it still complains something wrong with the running_mean size:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "syncBN.py", line 29, in &lt;module&gt;
    output = module(input)
  File "/home/ml/farleylai/miniconda3/envs/sinet37/lib/python3.7/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/ml/farleylai/miniconda3/envs/sinet37/lib/python3.7/site-packages/torch/nn/modules/container.py", line 92, in forward
    input = module(input)
  File "/home/ml/farleylai/miniconda3/envs/sinet37/lib/python3.7/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/ml/farleylai/miniconda3/envs/sinet37/lib/python3.7/site-packages/torch/nn/modules/batchnorm.py", line 459, in forward
    exponential_average_factor, self.eps)
  File "/home/ml/farleylai/miniconda3/envs/sinet37/lib/python3.7/site-packages/torch/nn/functional.py", line 1670, in batch_norm
    training, momentum, eps, torch.backends.cudnn.enabled
RuntimeError: running_mean should contain 2 elements not 100
</code></pre>
<blockquote>
<p>This would be a nice addition, we would be happy to merge a PR that adds this feature!</p>
</blockquote>
<p>So you are suggesting it is not intentional but something that could be completed?<br>
If that is the case, it seems like creating an issue on the GitHub repo makes more sense and I will look into the details under the hood.</p>
<p>Therefore to sum up, there is likely no luck for those with <code>BatchNorm1d</code> already to be converted to SyncBN transparently by default.</p>

username: albanD
cooked: <blockquote>
<p>The example input (2, 20) already contains a batch dim, indicating a batch of two 1D examples.</p>
</blockquote>
<p>That is not how batchnorm 1d works. Batchnorm 1d assumes an optional first batch dimension, then a channel dimension then an actual dimension. So the input is 2d without batch and 3d with batch.</p>
<blockquote>
<p>BTW, I tried to make it of size (1, 2, 20) but it still complains something wrong with the running_mean size:</p>
</blockquote>
<p>This is because you define your batchnorm as having <code>100</code> channels, but what you give as input has <code>2</code>.</p>

username: bonzogondo
cooked: <p>I am expiriencing the same problem as you <a class="mention" href="/u/farleylai">@farleylai</a> exactly now.</p>
<p>I am trying to run a model with ResNet backbone, which has only BatchNorm2d and a head network that have exactly ONE BatchNorm1d and that is exactly what causes problem.</p>
<p>The input to the BatchNorm1d in the forward function of the model is [64,2048].<br>
As suggested by <a class="mention" href="/u/alband">@albanD</a> I unsqueezed it in the forward function) so that the input shape is now [64, 1, 2048]. Next module is a Linear classifier, so I squeezed the output of the BatchNorm1d to again have [64, 2048] input to Linear layer. This helped in the sense that the forward pass is working, but in the backward pass I am getting now an error:</p>
<blockquote>
<p>RuntimeError: Function SyncBatchNormBackward returned an invalid gradient at index 1 - got [1] but expected shape compatible with [2048]</p>
</blockquote>
<p>Any suggestions <a class="mention" href="/u/alband">@albanD</a> ?</p>

username: albanD
cooked: <p>Do you have a small code sample so that I can reproduce that locally?</p>

username: bonzogondo
cooked: <p>Not really sure what you mean by ‘small code sample’. So I will try:</p>
<pre><code class="lang-auto">class Net(nn.Module):
    in_planes = 2048

    def __init__(self, num_classes, model_path,  model_name):
        super(Net, self).__init__()
        self.base = ResNet(block=Bottleneck,
                           layers=[3, 4, 6, 3])

        self.gap = nn.AdaptiveAvgPool2d(1)
        self.num_classes = num_classes
        self.bottleneck = nn.BatchNorm1d(self.in_planes)
        self.bottleneck.bias.requires_grad_(False)  # no shift
        self.classifier = nn.Linear(self.in_planes, self.num_classes, bias=False)

    def forward(self, x):

        global_feat = self.gap(self.base(x))  # (b, 2048, 1, 1)
        global_feat = global_feat.view(global_feat.shape[0], -1)  # flatten to (bs, 2048)

        feat = self.bottleneck(global_feat.unsqueeze(1))  ### To allow SyncBatchnorm

        cls_score = self.classifier(feat.squeeze())   ### To adjust for Linear layer input

        return cls_score, global_feat 


def train():
        model.train()
        optimizer.zero_grad()
        img, target = batch
        img = img.to(device)
        target = target.to(device)t
        score, feat = model(img)
        LOSS = loss_fn(score, feat, target)
        LOSS.backward()
        optimizer.step()
</code></pre>
<p>I shortened the code as much as possible to get the most important parts I think. The img size (input in train function) is [batch_size, 3, W, H] so standard for images.</p>
<p>EDIT: format and making a little clearer code</p>

username: bonzogondo
cooked: <p>I think I got it right now.</p>
<p>According to docs <a href="https://pytorch.org/docs/stable/nn.html#batchnorm1d" class="inline-onebox" rel="noopener nofollow ugc">torch.nn — PyTorch 2.1 documentation</a> and quoting:</p>
<blockquote>
<p>Parameters</p>
<ul>
<li><strong>num_features</strong> – C from an expected input of size (N,C,L) or L from input of size (N,L)</li>
</ul>
</blockquote>
<p>So I figured out that as we need 3d input and BatchNorm1d uses C as the num_features in three dimensional input, the singular dimension should be the last one.<br>
So instead of</p>
<p><code>feat = self.bottleneck(global_feat.unsqueeze(1))   # Which gives [bs, 1, 2048] </code><br>
I just did:<br>
<code>feat = self.bottleneck(global_feat.unsqueeze(-1))   # Which gives [bs, 2048, 1]</code></p>
<p>No more errors and training seems to run smoothly with SyncBatchNorm as well. Hope, this helps someone.</p>

username: farleylai
cooked: <blockquote>
<p>That is not how batchnorm 1d works. Batchnorm 1d assumes an optional first batch dimension, then a channel dimension then an actual dimension. So the input is 2d without batch and 3d with batch.</p>
</blockquote>
<p>As defined by the <code>BatchNorm1d</code>, the Input is expected to be of size <code>(N, L)</code> or <code>(N, C, L)</code> with batch dim first. What is optional is the additional channel dimension for <code>BatchNorm1d</code> from the documentation.</p>
<blockquote>
<p>This is because you define your batchnorm as having  <code>100</code>  channels, but what you give as input has  <code>2</code> .</p>
</blockquote>
<p><code>(1, 2, 20)</code> is due to the suggestion <code>adding .unsqueeze(0) to your input</code> but the resulting shape is not originally intended. By definition, whether the <code>100</code> is C or L in the previous example, <code>BatchNorm1d</code> produces the same results given <code>(N, 100)</code> or <code>(N, 100, 1)</code>. <code>(2, 100)</code> is already a batch input with <code>2</code> 1D features and matches the input accepted by <code>BatchNorm1d</code>. This has to be on the same page.</p>
<p>Now, get back to the issue with <code>SyncBatchNorm</code> conversion. Two questions:</p>
<ol>
<li>Does <code>SyncBatchNorm</code> wrapped <code>BatchNorm1d</code> behave as expected as before the conversion?</li>
</ol>
<p>The original <code>BatchNorm1d</code> takes both (N, L) or (N, C, L) and produces the same results as the following revised code segment shows. However, after converted to <code>SyncBatchNorm</code> which CHANGES the interface to ONLY accepts input of size (N, C, L). This conversion unlikely works transparently with existing models using <code>BatchNorm1d</code> to accept input of size <code>(N, L)</code>.</p>
<pre data-code-wrap="py"><code class="lang-py">import os
import copy
import torch
from torch import nn

with torch.no_grad():
    inputNL = torch.randn(2, 20).cuda()
    module = torch.nn.Sequential(
               torch.nn.Linear(20, 100),
               torch.nn.BatchNorm1d(100)
             ).cuda()
    moduleC = copy.deepcopy(module).cuda()
    moduleL = copy.deepcopy(module).cuda()
    moduleC.eval()
    moduleL.eval()

    # XXX: BatchNorm1d accepts (N, C, L) 
    outputNL = moduleC[0](inputNL)
    outputNCL = moduleC[1](outputNL.unsqueeze(-1))
    print('BatchNorm1d NCL:', outputNCL.shape, round(outputNCL.mean().item(), 7))

    # XXX: BatchNorm1d accepts (N, L) too
    outputNL = moduleL[0](inputNL)
    outputNL = moduleL[1](outputNL)
    print('BatchNorm1d NL:', outputNL.shape, round(outputNL.mean().item(), 7))

    os.environ['RANK'] = '0'
    os.environ['WORLD_SIZE'] = '1'
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '25791'
    torch.distributed.init_process_group(backend='nccl')

    moduleC = copy.deepcopy(module)
    moduleL = copy.deepcopy(module)
    moduleC = nn.SyncBatchNorm.convert_sync_batchnorm(moduleC)
    moduleL = nn.SyncBatchNorm.convert_sync_batchnorm(moduleL)
    moduleC.eval()
    moduleL.eval()

    # XXX: converted BatchNorm1d ONLY accepts (N, C, L) 
    outputNL = moduleC[0](inputNL)
    outputNCL = moduleC[1](outputNL.unsqueeze(-1))
    print('SyncBatchNorm NCL:', outputNCL.shape, round(outputNCL.mean().item(), 7))

    # FIXME: Converted BatchNorm1d never accepts (N, L)
    outputNL = moduleL[0](inputNL)
    outputNL = moduleL[1](outputNL)
    print('SyncBatchNorm NL:', outputNL.shape, round(outputNL.mean().item(), 7))
</code></pre>
<p>Sample output:</p>
<pre data-code-wrap="sh"><code class="lang-sh">BatchNorm1d NCL: torch.Size([2, 100, 1]) 0.0683341
BatchNorm1d NL: torch.Size([2, 100]) 0.0683341
SyncBatchNorm NCL: torch.Size([2, 100, 1]) 0.0683341
Traceback (most recent call last):
  File "syncBN.py", line 45, in &lt;module&gt;
    outputNL = moduleL[1](outputNL)
  File "/home/ml/farleylai/miniconda3/envs/sinet37/lib/python3.7/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/ml/farleylai/miniconda3/envs/sinet37/lib/python3.7/site-packages/torch/nn/modules/batchnorm.py", line 429, in forward
    self._check_input_dim(input)
  File "/home/ml/farleylai/miniconda3/envs/sinet37/lib/python3.7/site-packages/torch/nn/modules/batchnorm.py", line 417, in _check_input_dim
    .format(input.dim()))
ValueError: expected at least 3D input (got 2D input)
</code></pre>
<ol start="2">
<li>If not, what is the justification or workaround that does not require changing the existing model to wrap?</li>
</ol>
<p>One workaround is to reshape/unsqueeze(-1) the immediate input of size (N, L) to (N, C=L, L=1) before the converted <code>BatchNorm1d</code> as demonstrated by <a class="mention" href="/u/bonzogondo">@bonzogondo</a>. Unfortunately, this may not be scalable if the uses of <code>BatchNorm1d</code> are all over the place in existing models. There is no reshape layers in PyTorch to automate the <code>unsqeeze</code>. An alternative could be to identify whether the <code>BatchNorm</code> to wrap is 1D or not so that the SyncBatchNorm._check_input_dim(…) checks the same criteria as <code>BatchNorm1d</code> as sketched in the following. There may be some other exceptions but the goal should be to wrap existing models transparently.</p>
<pre data-code-wrap="py"><code class="lang-py">class SyncBatchNorm(nn.SyncBatchNorm):
    def _check_input_dim(self, input):
        if self._1d:
            if input.dim() != 2 and input.dim() != 3:
                raise ValueError('expected 2D or 3D input (got {}D input)'
                                    .format(input.dim()))
        elif input.dim() &lt;= 2:
            raise ValueError('expected at least 3D input (got {}D input)'
                             .format(input.dim()))

    @classmethod
    def convert_sync_batchnorm(cls, module, process_group=None):
        ...
        if isinstance(module, nn.modules.batchnorm._BatchNorm):
            module_output = SyncBatchNorm(module.num_features,
                                              module.eps, module.momentum,
                                              module.affine,
                                              module.track_running_stats,
                                              process_group)
            module_output._1d = isinstance(module, nn.modules.batchnorm.BatchNorm1d)
            ...
</code></pre>

username: Kirayue
cooked: <p>I faced the same problem. Any update?</p>