Convert tf batchnorm to pytorch

tree33: <p>Hi, I’m trying to convert a tensorflow pretrained model to pytorch.</p>
<p>The network is quite simple, just with one 1dconv and one 1dbn layer</p>
<p>The output is the same after the 1d conv layer. However, I cannot replicate the performance/output after the batchnorm layer</p>
<p>More specifically, there seems to be a bug in the original tensorflow code.</p>
<p>The author seems to forget to add the runningmean/variance of the batchnorm layer in tensorflow to the ‘trainable variables’, so they are never really updated in the model he provided. And the is_train flag is always set to True even in evaluation.</p>
<p>However, when I load the model in pytorch, then set the runningmean to all zeros and running variance to all ones, the output is different nomatter I set model to train or eval.</p>
<p>Nonetheless, the overall performance in the tensorflow model is great. Even if I remove all the bn layers, I can get the same output in my pytorch code, and when I add batchnorm back and retrain the network in pytorch, the performance is bad, about 20 ~ 30 percent lower,</p>
<p>Please check this original code  in tensorflow for details</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/doubleQ2018/Deep-learning-for-contact_map_v2/blob/master/libs/nets/network.py" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/doubleQ2018/Deep-learning-for-contact_map_v2/blob/master/libs/nets/network.py" target="_blank" rel="nofollow noopener">doubleQ2018/Deep-learning-for-contact_map_v2/blob/master/libs/nets/network.py</a></h4>
<pre><code class="lang-py">#!/usr/bin/env python

import tensorflow as tf
import numpy as np

def weight_variable(shape, regularizer, name="W"):
    if regularizer == None:
        initial = tf.truncated_normal(shape, stddev=0.1)
        return tf.Variable(initial, name)
    else:
        return tf.get_variable(name, shape, 
                initializer=tf.random_normal_initializer(), regularizer=regularizer)


def bias_variable(shape, name="b"):
    initial = tf.constant(0.1, shape=shape)
    return tf.Variable(initial, name)

### Incoming shape (batch_size, L(seqLen), feature_num)
### Output[:, i, j, :] = incoming[:. i, :] + incoming[:, j, :] + incoming[:, (i+j)/2, :]
</code></pre>

  This file has been truncated. <a href="https://github.com/doubleQ2018/Deep-learning-for-contact_map_v2/blob/master/libs/nets/network.py" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>And here’s my code for the problem above<br>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/YamamotoSayaka/form/blob/master/cvt_bn.py" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/YamamotoSayaka/form/blob/master/cvt_bn.py" target="_blank" rel="nofollow noopener">YamamotoSayaka/form/blob/master/cvt_bn.py</a></h4>
<pre><code class="lang-py">import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np
import tensorflow as tf
import pickle

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

filter_size = 17
in_channels = 26
out_channels = 28
inp = np.random.rand(1,100,26)


# tensorflow model
def weight_variable(shape, name='W'):
    initial = tf.truncated_normal(shape, stddev=0.1)
</code></pre>

  This file has been truncated. <a href="https://github.com/YamamotoSayaka/form/blob/master/cvt_bn.py" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>

ptrblck: <aside class="quote no-group" data-username="tree33" data-post="1" data-topic="50565">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tree33/48/10609_2.png" class="avatar"> tree33:</div>
<blockquote>
<p>The author seems to forget to add the runningmean/variance of the batchnorm layer in tensorflow to the ‘trainable variables’, so they are never really updated in the model he provided. And the is_train flag is always set to True even in evaluation.</p>
</blockquote>
</aside>
<p>I’m not that familiar with tensorflow, but are the running estimates treated as “trainable variables”? Since they are updated based on the current batch statistic, they won’t be updated using the gradient + optimizer.<br>
Are you sure the running estimates are not updated, if they are not in these “trainable variables” list?</p>
<aside class="quote no-group" data-username="tree33" data-post="1" data-topic="50565">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tree33/48/10609_2.png" class="avatar"> tree33:</div>
<blockquote>
<p>However, when I load the model in pytorch, then set the runningmean to all zeros and running variance to all ones, the output is different nomatter I set model to train or eval.</p>
</blockquote>
</aside>
<p>In your current code you are setting the <code>weight</code> and <code>bias</code> to these values: <a href="https://github.com/YamamotoSayaka/form/blob/acdee69ff2a87cf511d64c37f2a41597adba0ab7/cvt_bn.py#L77-L78">lines of code</a>.<br>
These parameters correspond to the gamma and beta parameters from the original batchnorm paper.<br>
However, the <code>running_mean</code> and <code>running_var</code> will be initialized by default to zeroes and ones, respectively.</p>

tree33: <p><a class="mention" href="/u/ptrblck">@ptrblck</a><br>
Well I’m also not quite familiar with tensorflow myself,  and that’s the reason why I want to convert this model to pytorch in the first place.</p>
<p>In fact after looking into tensorflow’s doc, I think the running mean and variance<br>
are not updated in the github’s repo, and they are not in the ‘trainable variables’ by default in tensorflow. These lines of code are missing.</p>
<pre><code class="lang-auto">  update_ops = tf.compat.v1.get_collection(tf.GraphKeys.UPDATE_OPS)
  with tf.control_dependencies(update_ops):
    train_op = optimizer.minimize(loss)
</code></pre>
<p>What’s more , I load the pretrained model provided by the author in tensorflow, and it gives me a reasonable performance. And in fact, I can print these weights of the first batch norm layer with the following lines of code:</p>
<pre><code class="lang-auto">bn_weight = tf.get_collection(key=tf.GraphKeys.GLOBAL_VARIABLES, scope="ResidualBlock_1D_0/BatchNorm/    ")
sess.run(bn_weight)
</code></pre>
<p>And here’s the output</p>
<pre><code class="lang-auto">(Pdb) print(bn_weight)
[&lt;tf.Variable 'ResidualBlock_1D_0/BatchNorm/beta:0' shape=(28,) dtype=float32_ref&gt;, &lt;tf.Variable 'ResidualBlock_1D_0/BatchNorm/moving_mean:0' shape=(28,) dtype=float32_ref&gt;, &lt;tf.Variable 'ResidualBlock_1D_0/BatchNorm/moving_variance:0' shape=(28,) dtype=float32_ref&gt;]
(Pdb) sess.run(bn_weight)
[array([ 0.00602053,  0.0050063 ,  0.00634659,  0.00489874, -0.00110139,
       -0.0022672 ,  0.00327282,  0.00230329, -0.00308617,  0.0024446 ,
        0.00504038,  0.00412193, -0.00564392,  0.00558108,  0.00340598,
        0.00164128, -0.00063161,  0.00081446, -0.00207535,  0.00168148,
       -0.00086416,  0.00218814,  0.00166753, -0.00242493,  0.00044503,
       -0.00097768,  0.00121772,  0.00186761], dtype=float32), array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32), array([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
       1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], dtype=float32)]
</code></pre>
<p>And I think I’m doing the right thing in pytorch by setting the mean to zeros and variance to ones. However, the output is just different, wired.</p>
<p>I also tried to set track_running_stats to False or affine to True in pytorch, no help at all. <img src="https://discuss.pytorch.org/images/emoji/apple/roll_eyes.png?v=9" title=":roll_eyes:" class="emoji" alt=":roll_eyes:"><img src="https://discuss.pytorch.org/images/emoji/apple/roll_eyes.png?v=9" title=":roll_eyes:" class="emoji" alt=":roll_eyes:"></p>

Alpha: <p>Hi, <a class="mention" href="/u/ptrblck">@ptrblck</a></p>
<p>A similar question about Jax and Pytorch.</p>
<p>For BatchNorm, I noticed that there is a difference between <a href="https://github.com/deepmind/dm-haiku/blob/main/haiku/_src/batch_norm.py#L188-L200" rel="noopener nofollow ugc">Jax</a> and <a href="https://github.com/pytorch/pytorch/blob/700109eb630b79fd65cb93becb7f2d14f93bdb5c/aten/src/ATen/native/Normalization.cpp#L238-L260" rel="noopener nofollow ugc">Pytorch</a>.</p>
<p>I.e.</p>
<p>Jax computation order:  <code>(inputs - mean) *  (scale * rsqrt(var + eps))</code></p>
<p>Pytorch:  <code>(inputs - mean) * scale * rsqrt(var + eps)</code>.</p>
<p>I know they are theoretically equivalent.</p>
<p>Do you think this computation order would make difference?</p>
<p>Thanks in advance.</p>

ptrblck: <p>You might be running into rounding errors using floating point numbers, which could then create a small error between both implementations, but I wouldn’t expect to see any difference in your training etc.</p>