Getting different outputs from batchnorm network with different batch sizes at test time

blackberry: <p>I trained a DenseNet to compute image embeddings of size 256, to match images.<br>
At test time, when I use different batch sizes (64, 256, 512), the network outputs (embedding) are slightly different (I do model.eval()).<br>
The difference seems to be proportional to the difference in the batch size.</p>
<p>I suspect, this is due to batchnorm (track_running_stats).</p>
<p><strong>Update:</strong> Indeed, it is related to batchnorm’s <code>track_running_stats</code>.</p>
<p>If I set <code>track_running_stats</code> to <code>False</code>, but use a large batch size (512), it works well (high matching accuracy). When I use small batch size (8), the accuracy drops a few percent.</p>
<p>If I leave <code>track_running_stats</code> as <code>True</code> (default), even though the outputs for different batch sizes differ slightly, the matching accuracy is consistent across batch sizes, and slightly lower than <code>track_running_stats = False</code> with large batch size.</p>
<p>How to solve this?<br>
I should get the same output regardless of the batch size.<br>
Is it possible with a network with batchnorm?</p>
<p>PyTorch v1.4</p>
<p><strong>Update 2: Typical Use case Scenario:</strong></p>
<ul>
<li>I train a batchnorm CNN network with a single embedding layer of size 256 with metric loss, and large batch size (256).</li>
<li>I also validate/test using a large batch size (512)</li>
<li>Got a good model, let me <strong>deploy</strong> it. Compute the embeddings for <strong>millions of images</strong> in the database using a l<strong>arge batch size</strong> (512). Great!</li>
<li>At <strong>test time</strong>, though, the network receives a <strong>single image</strong> at a time, in real-time! Oooops! The batch size is now 1. Even if I feed one of the images from the database, I do not get the same embeddings! If I match them, I do not get a distance of 0! (We cannot explain this to an outsider).</li>
<li>Hence, the problem.</li>
</ul>

chetan_patil: <p>Hi,<br>
do you use <code>model.eval()</code>  before you pass the inputs to the network ?</p>

blackberry: <p>Yes, as I said in my post.</p>
<p>The problem is related to the batch statistics. I would like to see if there is a good solution that works.</p>

ptrblck: <aside class="quote no-group" data-username="blackberry" data-post="1" data-topic="88622">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/b/ecd19e/48.png" class="avatar"> blackberry:</div>
<blockquote>
<p>I should get the same output regardless of the batch size.<br>
Is it possible with a network with batchnorm?</p>
</blockquote>
</aside>
<p>This won’t be possible with <code>track_running_stats=False</code>, since the batchnorm layers will calculate the current batch statistics, which will vary based on the batch size.<br>
However, <code>track_running_stats=True</code> should give you the same outputs (up to floating point precision, if non-deterministic methods were used or if you haven’t used the setup from the <a href="https://pytorch.org/docs/stable/notes/randomness.html">reproducibility docs</a>).<br>
How large is the absolute difference with <code>track_running_stats=True</code>?</p>

blackberry: <p>Thanks Peter.</p>
<p>I have updated my post with a typical use case scenario (also how I use it).</p>
<p>Indeed, based on my preliminary experiments, it is not possible to get the same embeddings with <code>track_running_stats=False</code>. Now, trying to find a solution with <code>track_running_stats=True</code>.</p>
<p>The distances between embeddings of the same images with different batch sizes are summarized <a href="https://discuss.pytorch.org/t/getting-different-outputs-from-batchnorm-network-with-different-batch-sizes-at-test-time/88622/7">below</a>.</p>
<p>I wish, it was possible to also save the running mean/std of batchnorm layers from training and use them at test time with a small momentum. This would mitigate the problem.</p>
<p>A workaround might be to run the model on the training set with large batch size, right before switching to test. I have not tried this yet.</p>
<p>Do you have a better idea?</p>

ptrblck: <aside class="quote no-group" data-username="blackberry" data-post="5" data-topic="88622">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/b/ecd19e/48.png" class="avatar"> blackberry:</div>
<blockquote>
<p>I wish, it was possible to also save the running mean/std of batchnorm layers from training and use them at test time with a small momentum. This would mitigate the problem.</p>
</blockquote>
</aside>
<p>If you are loading the <code>state_dict</code>, the running stats will also be restored.<br>
Since the validation and test cases are executed via <code>model.eval()</code>, the <code>momentum</code> term won’t be used. Instead the running stats will be used directly to normalize the data.</p>
<p>Could you post the model architecture, so that we could have a look?</p>

blackberry: <p>If the running stats are saved and re-stored and used at test time (the docs also say so), then there should not be any difference, regardless of the batch size. I use <code>torch.save('model.pkl', model)</code> and <code>model = torch.load('model.pkl')</code>.</p>
<p>Here is the model, which is just a DenseNet:</p>
<pre><code class="lang-auto">class DenseNet(nn.Module):

   def __init__(self, layers=169, pretrained=True, emb_size=256):
       super(DenseNet, self).__init__()

       if layers == 121:
           self.model = models.densenet121(pretrained=pretrained)
           fc_in = 1024
       elif layers == 161:
           self.model = models.densenet161(pretrained=pretrained)
           fc_in = 2208
       elif layers == 169:
           self.model = models.densenet169(pretrained=pretrained)
           fc_in = 1664
       elif layers == 201:
           self.model = models.densenet201(pretrained=pretrained)
           fc_in = 1920

       self.model.classifier = nn.Linear(fc_in, emb_size, bias=False)


   def forward(self, x, norm=True):
       x = self.model.features(x)
       x = F.relu(x, inplace=True)
       x = F.adaptive_avg_pool2d(x, (1, 1))
       x = x.view(x.size(0), -1)

       x = self.model.classifier(x)
       if norm:
           x = F.normalize(x)

       return x
</code></pre>
<p>I used it like this:</p>
<pre><code class="lang-auto">model = DenseNet(layers=169, pretrained=True, emb_size=256)
model = model.cuda()
model.train()
torch.set_grad_enabled(True)
# train
torch.save('model.pkl', model)

# Test time
model = torch.load('model.pkl')
model.eval()
torch.set_grad_enabled(False)
# do test
</code></pre>
<p>I suppose <code>model.train()/model.eval()</code> on the DenseNet object recursively calls train/eval on all the modules it contains. If it is not the case, then that is the problem [update: no, model.train()/eval() works as expected].</p>
<p><strong>Update:</strong> Here are some numbers, L1 distances between the DenseNet 169 embeddings of the same 5 images (out of 512) using different batch sizes in model.eval() mode.</p>
<pre><code class="lang-auto"># load the model, switch to eval(), run it on 512 images with different batch sizes, and compute L1 distances in memory, without saving to disk.
# batch sizes: 512 vs 1 - five runs, seems to be non-deterministic
tensor([1.2099e-05, 1.2104e-05, 5.7479e-06, 1.1539e-05, 9.9048e-06])
tensor([1.2099e-05, 1.2104e-05, 5.7479e-06, 1.1539e-05, 9.9048e-06])
tensor([1.4106e-05, 1.2881e-05, 6.3117e-06, 1.1826e-05, 1.0850e-05])
tensor([1.4106e-05, 1.2881e-05, 6.3117e-06, 1.1826e-05, 1.0850e-05])
tensor([1.2099e-05, 1.2104e-05, 5.7479e-06, 1.1539e-05, 9.9048e-06])

# batch sizes: 512 vs 1
# Forward 10 batches  of 512 in train mode, then switch to eval mode
tensor([1.3889e-05, 1.3013e-05, 7.1302e-06, 1.2253e-05, 1.3526e-05])

# batch sizes: 512 vs 4
tensor([1.4723e-05, 1.4445e-05, 8.1474e-06, 1.2522e-05, 1.2616e-05])

# batch sizes: 512 vs 64
tensor([1.4181e-05, 1.3389e-05, 6.6453e-06, 1.1660e-05, 1.1504e-05])

# batch sizes: 512 vs 256
tensor([3.9932e-06, 4.5871e-06, 2.6375e-06, 3.3856e-06, 4.4151e-06])

# batch sizes: 512 vs 512
tensor([0., 0., 0., 0., 0.])

# ImageNet Pre-trained DenseNet 169 with new non-trained embedding layer
# batch sizes: 512 vs 1
tensor([1.6982e-05, 1.5111e-05, 1.6363e-05, 1.5956e-05, 1.8549e-05])

# Random non-trained DenseNet
# batch sizes: 512 vs 1
tensor([4.8725e-06, 4.9665e-06, 4.8112e-06, 4.9490e-06, 4.7400e-06])

# ImageNet Pre-trained ResNet50 with new non-trained embedding layer
# batch sizes: 512 vs 1
tensor([1.7733e-05, 1.3676e-05, 1.5779e-05, 1.1592e-05, 1.1086e-05])
</code></pre>
<p>Based on these, it seems to be changing with batch size, and hence most probably related to batchnorm. PyTorch’s batchnorm layer uses an epsilon value of <code>1e-5</code>, and the L1 distances are in the same range. Maybe related to this?</p>

ptrblck: <p>These differences are most likely created by the limited floating point precision, not by batchnorm or any other layers.<br>
Usually you would expect a relative error of <code>~1e-6</code>, but if you are using a deep model, these errors might accumulate.</p>

blackberry: <p>Initially, I also thought so. After the above experiments, I think there is a problem somewhere (probably due to floating point precision), and  most probably the problem is in the batchnorm layer, since the distance is proportional to the difference in batch size, and the distance is zero when the batch sizes are the same (see 512 vs 512).</p>

ptrblck: <p>How did you narrow down the batchnorm layers? Did you replace them and got a zero absolute error or what’s the reasoning?<br>
Due to the limited floating point precision operations such as e.g. <code>sum</code> will also suffer from this effect:</p>
<pre><code class="lang-python">x = torch.randn(100, 100, 100)
sum1 = x.sum()
sum2 = x.sum(0).sum(0).sum(0)
print((sum1 - sum2).abs())
&gt; tensor(0.0016)
</code></pre>

blackberry: <p>No, I have not done that. That is probably the best way to show it.</p>
<p>But, as I said: (1) it changes with batch size, (2) I get 0 distance if I use the same batch size.</p>
<p>In your example, you are using different sequence of operations, in contrast to my example, in which the network is supposed to be the same (same sequence of operations).</p>