Batchnorm for different sized samples in batch

Dan_Erez: <p>I use batchnorm 1d on batches which are padded to the max length of the samples. It dawned on me that batch norm isn’t fed a mask so it has no way of knowing which are valid timesteps in each sequence. Wouldn’t this mess with batch norm? And more importantly wouldn’t it be very different if I change batch size? Is there a way around this?</p>

gollum: <p>Hi <a class="mention" href="/u/dan_erez">@Dan_Erez</a>! Did you find anything to solve the problem? I’m facing the same problem. Is there a solution similar to the one in tensorflow <a href="https://www.tensorflow.org/api_docs/python/tf/layers/BatchNormalization#input_mask" rel="nofollow noopener">here</a> in PyTorch?</p>

gollum: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Please, can you help me with this?<br>
Basically, I have a tensor padded with zeros in the end. If I feed this into <code>torch.nn.BatchNorm1d</code> it will consider those as well. I also have a mask (binary) for the padded tensor. Is there something in PyTorch to tackle this?</p>

ptrblck: <p>As far as I understand your use case, you are creating a batch (3-dimensional: <code>[batch_size, channels, seq_len]</code>), where some tensors were zero-padded in the last dimension.<br>
Is that correct?<br>
Now you would like to ignore the padded inputs in the batchnorm layer, i.e. not being taken into account for the running stats or what would the desired behavior be?</p>

gollum: <p>Thanks for the reply! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
Yes precisely. This is exactly what I need. If the zeros are taken into account, it will be wrong. Need to ignore those.</p>

ptrblck: <p>I’m not aware of any built-in method, so you might need to implement it manually.</p>
<p>Maybe you could use <a href="https://github.com/ptrblck/pytorch_misc/blob/20e8ea93bd458b88f921a87e2d4001a4eb753a02/batch_norm_manual.py#L60" rel="nofollow noopener">this manual example of the batch norm calculation</a> as a starter and change the <code>mean</code> and <code>var</code> calculation using the masked method:</p>
<pre><code class="lang-python"># Create dummy input
x = torch.randn(2, 3, 10)
x[0, 0, 5:] = 0.
x[0, 1, 6:] = 0.
x[0, 2, 7:] = 0.
x[1, 0, 8:] = 0.
x[1, 1, 9:] = 0.

# Use mask for manual calculation
mask = x!=0
mask_mean = (x.sum(2) / mask.float().sum(2)).mean(0)

# Alternatively rescale BatchNorm1d.running_mean
mean = x.mean([0, 2])
mean * x.size(2) / (x.size(2) - (x==0).float().sum(2)).sum()
</code></pre>
<p>The second example would work, if you would like to use the PyTorch batch norm implementation (e.g. for performance reasons) and “rescale” the running estimates.</p>
<p>Let me know, if that helps.</p>

Dan_Erez: <p>thanks for the reply. can you explicity show how you would operate on the batch norm parameter?</p>

isarandi: <p>I wrote a solution to do this fast, explained as comments in the code. Let me know if you find any bugs.</p>
<pre><code class="lang-auto">    def masked_batchnorm1d_forward(x, mask, bn):
        """x is the input tensor of shape [batch_size, n_channels, time_length]
            mask is of shape [batch_size, 1, time_length]
            bn is a BatchNorm1d object
        """
        if not self.training:
            return bn(x)

        # In each example of the batch, we can have a different number of masked elements
        # along the time axis. It would have to be represented as a jagged array.
        # However, notice that the batch and time axes are handled the same in BatchNorm1d.
        # This means, we can merge the time axis into the batch axis, and feed BatchNorm1d
        # with a tensor of shape [n_valid_timesteps_in_whole_batch, n_channels, 1],
        # as if the time axis had length 1.
        #
        # So the plan is:
        #  1. Move the time axis next to the batch axis to the second place
        #  2. Merge the batch and time axes using reshape
        #  3. Create a dummy time axis at the end with size 1.
        #  4. Select the valid time steps using the mask
        #  5. Apply BatchNorm1d to the valid time steps
        #  6. Scatter the resulting values to the corresponding positions
        #  7. Unmerge the batch and time axes
        #  8. Move the time axis to the end 

        n_feature_channels = x.shape[1]
        time_length = x.shape[2]
        reshaped = x.permute(0, 2, 1).reshape(-1, n_feature_channels, 1)
        reshaped_mask = mask.reshape(-1, 1, 1) &gt; 0
        selected = torch.masked_select(reshaped, reshaped_mask).reshape(-1, n_feature_channels, 1)
        batchnormed = bn(selected)
        scattered = reshaped.masked_scatter(reshaped_mask, batchnormed)
        backshaped = scattered.reshape(-1, time_length, n_feature_channels).permute(0,2,1)
        return backshaped
</code></pre>

Bubas3000: <p>Batchnorm1d has learnable parameters. Is there any problem with having different sized batches?<br>
Thank you</p>