Loading TensorFlow GRUCell weights into PyTorch

vadimkantorov: <p>The checkpoint for Slot Attention object discovery model: <a href="https://github.com/google-research/google-research/tree/master/slot_attention#pre-trained-model-checkpoints" class="inline-onebox" rel="noopener nofollow ugc">google-research/slot_attention at master · google-research/google-research · GitHub</a> (<a href="https://console.cloud.google.com/storage/browser/gresearch/slot-attention" rel="noopener nofollow ugc">https://console.cloud.google.com/storage/browser/gresearch/slot-attention</a>) has a <a href="https://github.com/google-research/google-research/blob/master/slot_attention/model.py#L65" rel="noopener nofollow ugc">tensorflow.keras.layers.GRUCell</a> (from TFv2).</p>
<p>Corresponding keys (with shapes) in the checkpoint are:</p>
<pre><code class="lang-python">{
  'network/layer_with_weights-0/slot_attention/gru/recurrent_kernel/.ATTRIBUTES/VARIABLE_VALUE' : torch.Size([64, 192]), 
  'network/layer_with_weights-0/slot_attention/gru/kernel/.ATTRIBUTES/VARIABLE_VALUE' : torch.Size([64, 192]),
  'network/layer_with_weights-0/slot_attention/gru/bias/.ATTRIBUTES/VARIABLE_VALUE' : torch.Size([2, 192])
}
</code></pre>
<p>Parameters for PyTorch one seem to be</p>
<pre><code class="lang-python">{
  'slot_attention.gru.weight_ih': torch.Size([192, 64]),
  'slot_attention.gru.weight_hh': torch.Size([192, 64]),
  'slot_attention.gru.bias_ih': torch.Size([192]), 
  'slot_attention.gru.bias_hh': torch.Size([192])
}
</code></pre>
<p>How does one convert them to PyTorch GRUCell (especially in the view of<br>
<a href="https://discuss.pytorch.org/t/grucell-is-different-in-pytorch-and-tensorflow/6630" class="inline-onebox">GRUcell is different in pytorch and tensorflow</a>)?</p>
<p>Would the following be correct (in pseudo-code)?:</p>
<pre><code class="lang-python">weight_hh = recurrent_kernel.t()
weight_ih = kernel.t()
bias_ih, bias_hh = bias.unbind()
</code></pre>
<p>What should be the bias format? (maybe <a href="https://github.com/tensorflow/tensorflow/issues/44948#issuecomment-735047539" rel="noopener nofollow ugc">related</a>)</p>
<p>Thank you!</p>

JamesWright: <p>Hey <a class="mention" href="/u/vadimkantorov">@vadimkantorov</a> , were you able to figure this out? I am facing exactly the same issue. Thanks.</p>

vadimkantorov: <p>I don’t remember for sure now, but I think my snippet above is correct. Best would be to test against TF outputs, but I don’t have the setup…</p>

JamesWright: <p>I’ve tested it against TensorFlow/Keras outputs, it does not work (output is different). I’ve found following conversion to work for <strong>GRU</strong> layer (<em>TensorFlow r2.4rc0</em> - <code>tf.keras.layers.GRU</code> to <em>PyTorch 1.9.0</em> - <code>torch.nn.GRU</code>):</p>
<ol>
<li>Save weights from TensorFlow model (for example to <code>.npz</code> file):</li>
</ol>
<pre><code class="lang-python">import random
import numpy as np
import tensorflow as tf

SEED=1995

np.random.seed(SEED)
random.seed(SEED)
tf.random.set_seed(SEED)

gru = tf.keras.layers.GRU(
    units=5,
    return_sequences=True,
    kernel_initializer=tf.keras.initializers.GlorotUniform(seed=SEED),
    recurrent_initializer=tf.keras.initializers.Orthogonal(seed=SEED),
    bias_initializer=tf.keras.initializers.Zeros
)

y_tf = gru(tf.ones((1, 3, 5)), training=False)  # forward pass with ones

np.savez(
    'tf_model_weights.npz', 
    gru_kernel=gru.weights[0].numpy(), 
    gru_recurrent_kernel=gru.weights[1].numpy(),
    gru_bias=gru.weights[2].numpy()
)
</code></pre>
<ol start="2">
<li>Load weights into PyTorch model:</li>
</ol>
<pre><code class="lang-python">import random as r
import numpy as np
import torch

SEED=1995
torch.set_printoptions(precision=8)

r.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)

from speechbrain.nnet.RNN import GRU, LSTM

npz_weights = np.load('tf_model_weights.npz')


def convert_input_kernel(kernel):
    kernel_r, kernel_z, kernel_h = np.hsplit(kernel, 3)
    return np.concatenate((kernel_z.T, kernel_r.T, kernel_h.T))
    

def convert_recurrent_kernel(kernel):
    kernel_r, kernel_z, kernel_h = np.hsplit(kernel, 3)
    return np.concatenate((kernel_z.T, kernel_r.T, kernel_h.T))


def convert_bias(bias):
    bias = bias.reshape(2, 3, -1) 
    return bias[:, [1, 0, 2], :].reshape((2, -1))


gru = torch.nn.GRU(
    hidden_size=5,
    input_size=5,
    num_layers=1,
    bidirectional=False,
    batch_first=True
)
for pn, p in gru.named_parameters():
    if 'weight_ih' in pn:
        p.data = torch.from_numpy(convert_input_kernel(npz_weights['gru_kernel']))
    elif 'weight_hh' in pn:
        p.data = torch.from_numpy(convert_recurrent_kernel(npz_weights['gru_recurrent_kernel']))
    elif 'bias_ih' in pn:
        p.data = torch.from_numpy(convert_bias(npz_weights['gru_bias'])[0])
    else:
        p.data = torch.from_numpy(convert_bias(npz_weights['gru_bias'])[1])
</code></pre>
<ol start="3">
<li>Test output:</li>
</ol>
<p><strong>TensorFlow:</strong></p>
<pre><code class="lang-python">&gt;&gt;&gt; y_tf
    &lt;tf.Tensor: shape=(1, 3, 5), dtype=float32, numpy=
    array([[[-0.36656722, -0.4693069 , -0.16722648,  0.36081928, 0.1643753 ],
            [-0.4628504 , -0.6815055 , -0.18605384,  0.58125013, 0.2494137 ],
            [-0.48067108, -0.7698146 , -0.16238967,  0.70518744, 0.3005259 ]]], dtype=float32)&gt;
</code></pre>
<p><strong>PyTorch:</strong></p>
<pre><code class="lang-auto">&gt;&gt;&gt; gru.eval()
    GRU(5, 5, batch_first=True)
&gt;&gt;&gt; y_pt, _ = gru(torch.ones(1, 3, 5))
&gt;&gt;&gt; y_pt
    tensor([[[-0.36656719, -0.46930692, -0.16722649,  0.36081928,  0.16437532],
             [-0.46285039, -0.68150550, -0.18605389,  0.58125013,  0.24941370],
             [-0.48067111, -0.76981461, -0.16238970,  0.70518738,  0.30052590]]],
           grad_fn=&lt;TransposeBackward1&gt;)
</code></pre>
<p>Hope someone will find this helpful.</p>

Yann: <p>Hi, I am trying to do the opposite: tranferring from Pytorch to TensorFlow, I tried to revert your code as follows:</p>
<pre><code class="lang-auto">def convert_kernel(kernel):
    kernel_z, kernel_r, kernel_h = np.vsplit(kernel, 3)
    return np.concatenate((kernel_r.T, kernel_z.T, kernel_h.T), axis=1)

def convert_bias(bias):
    bias = bias.reshape(2, 3, -1) 
    return bias[:, [1, 0, 2], :].reshape((2, -1))

# GRU Layer:

kernel_input = convert_kernel(torch_params[6])
kernel_h = convert_kernel(torch_params[7])
bias = convert_bias(np.stack((torch_params[8], torch_params[9]), axis=0))

model_keras.layers[8].set_weights([kernel_input, 
                                   kernel_h, 
                                   bias])
</code></pre>
<p>… but this doesn’t work. Would you have an idea why please? <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>Thanks!</p>

Yann: <p>Nevermind, it works, the issue was right before that (TensorFlow Reshape not behaving the same as Pytorch view() apparently)</p>