TF 1.x batch_normalization to PyTorch BatchNorm2d

fabricio: <p>I’m trying to convert an old TensorFlow 1.x model to PyTorch. I have created a script to export the TF model weights and load them to the PyTorch model. However, I still have problems with the batch normalization. This is my old source code:</p>
<pre><code class="lang-auto">def cnn_block(inputs, filters, kernel_size):
    cnn = tf.layers.conv2d(
        inputs=inputs,
        filters=filters,
        kernel_size=kernel_size,
        padding="same",
        activation=None,
        use_bias=False
    )
    output = tf.layers.batch_normalization(
        inputs=cnn,
        momentum=0.9,
        epsilon=1e-5,
        center=True,
        scale=True,
    )
    return output
</code></pre>
<p>I have implemented the following PyTorch code:</p>
<pre><code class="lang-auto">import torch
from torch import nn

class CNNBlock(nn.Module):

    def __init__(self, in_channels: int, out_channels: int, kernel_size: int):
        super().__init__()

        self.model = nn.Sequential(
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=kernel_size,
                padding=kernel_size // 2,
                bias=False
            ),
            nn.BatchNorm2d(
                num_features=out_channels,
                eps=1e-5,
                momentum=0.9,
                affine=True
            )
        )

    def forward(self, tensor: torch.Tensor) -&gt; torch.Tensor:
        return self.model(tensor)
</code></pre>
<p>However, the <code>nn.BatchNorm2d()</code> method returns different results than <code>tf.layers.batch_normalization()</code> method.</p>
<p>In the load weight routine, I copy the <code>gamma</code> value to <code>model.weight</code>, <code>beta</code> value to <code>model.bias</code>, <code>moving_mean</code> value to <code>model.running_mean</code>, and <code>moving_variance</code> value to <code>model.running_var</code>. However, the results are quite differents.</p>
<p>If I comment the batch normalization in both source codes, the <code>Conv2d()</code> returns basically the same value.</p>

ptrblck: <p>The <code>momentum</code> definition in both frameworks differ.<br>
While TF uses a <a href="https://www.tensorflow.org/api_docs/python/tf/keras/layers/BatchNormalization">default of 0.99</a> as it’s multiplied with the moving stats:</p>
<blockquote>
<p>moving_mean = moving_mean * momentum + mean(batch) * (1 - momentum)<br>
moving_var = moving_var * momentum + var(batch) * (1 - momentum)</p>
</blockquote>
<p>PyTorch uses a <a href="https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html">default of 0.1</a> as it’s multiplied with the new observed value:</p>
<blockquote>
<p>This <code>momentum</code> argument is different from one used in optimizer classes and the conventional notion of momentum. Mathematically, the update rule for running statistics here is x^new=(1−momentum)×x^+momentum×xt​, where x^ is the estimated statistic and xt​ is the new observed value.</p>
</blockquote>

fabricio: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>, thanks for the fast reply.<br>
Therefore, must I implement a new <code>nn.Module</code> to normalize in the same way as in TF?</p>

ptrblck: <p>You could change the <code>momentum</code> via subtracting it from 1 first and check if the results would match.</p>

fabricio: <p>As I understood, the <code>momentum</code> is used only during the training, right?</p>
<p>I’m exporting a trained model from TF to PyTorch. Therefore, the <code>momentum</code> doesn’t have any difference in the normalization calculation.</p>
<p>In my script, I load the trained TF weights using <code>pickle</code> package, and copy the weights to the PyTorch model. However, when I send the same tensor to both models (TF and PyTorch), the normalization results differ. This is the code I’m using to load the trained weights:</p>
<pre><code class="lang-auto">data = pickle.load(FILENAME)

model = tf_model.layer[1].model[2]
tensor = torch.tensor(
    data["normalization_layer/gamma:0"].copy(),
    requires_grad=True, dtype=model.weight.dtype
)
model.weight.copy_(tensor)
tensor = torch.tensor(
    data["normalization_layer/beta:0"].copy(),
    requires_grad=True, dtype=model.bias.dtype
)
model.bias.copy_(tensor)
</code></pre>

fabricio: <p>I found a solution! The problem was when I exported the model to ONNX format. I must set the <code>TrainingMode</code> to <code>EVAL</code>. Both TensorFlow 1.x and PyTorch models calculate the same values (almost because there is a slight difference in the float numbers). This was my final solution:</p>
<pre><code class="lang-auto">data = pickle.load(FILENAME)

model = tf_model.layer[1].model[2]
tensor = torch.tensor(
    data["normalization_layer/gamma:0"].copy(),
    requires_grad=True, dtype=model.weight.dtype
)
model.weight.copy_(tensor)
tensor = torch.tensor(
    data["normalization_layer/beta:0"].copy(),
    requires_grad=True, dtype=model.bias.dtype
)
model.bias.copy_(tensor)
tensor = torch.tensor(
    data["normalization_layer/moving_mean:0"].copy(),
    requires_grad=True, dtype=model.running_mean.dtype
)
model.running_mean.copy_(tensor)
tensor = torch.tensor(
    data["normalization_layer/moving_variance:0"].copy(),
    requires_grad=True, dtype=model.running_var.dtype
)
model.running_var.copy_(tensor)

batch_size = 1
tensor = torch.randn(batch_size, C, H, W, requires_grad=True)

torch.onnx.export(
    model,
    tensor,
    f"model.onnx",
    export_params=True,
    opset_version=17,
    do_constant_folding=False,
    training=torch.onnx.TrainingMode.EVAL,
    input_names = ["input"],
    output_names = ["output"],
    dynamic_axes={
        "input" : {0 : "batch_size"},
        "output" : {0 : "batch_size"}
    }
)
</code></pre>