Implementation difference between TensorFlow v1 LSTMBlockFusedCell and PyTorch LSTM

username: kevs
cooked: <p>I am attempting to translate a tensorflow <a href="http://man.hubwiz.com/docset/TensorFlow.docset/Contents/Resources/Documents/api_docs/python/tf/contrib/rnn/LSTMBlockFusedCell.html" rel="noopener nofollow ugc">LSTMBlockFusedCell</a> model to pytorch <a href="https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html" rel="noopener nofollow ugc">LSTM</a>, but I’m not getting the same outputs with identical input and weights in both models. I believe this is due to how the weights are being set for the torch model; in the code snippet beneath the TensorFlow weight has the shape <code>(400, 164)</code> whilst the PyTorch weights has the shape <code>(400,64)</code> and <code>(400,100)</code> for <code>torch_lstm.weight_ih_l0</code> and <code>torch_lstm.weight_hh_l0</code> respectively. I addressed this inconsistency by using the first <code>64</code> elements as <code>weight_ih_l0</code> and the proceeding <code>100</code> elements as <code>weight_hh_l0</code>. According to <a href="https://www.adrian.idv.hk/2020-12-31-torch2tf/" rel="noopener nofollow ugc">this article</a>, TensorFlow uses right-multiplication instead of PyTorch left-multiplication which is why I need to transpose the weight. Also I am setting the bias to 0 (rendering it useless) for debugging.</p>
<pre><code class="lang-python">import tensorflow as tf
import numpy as np
import torch

# setup tensorflow LSTM
tf_lstm = tf.contrib.rnn.LSTMBlockFusedCell(num_units=100)
inp = tf.placeholder(tf.float32, shape=(1, 50, 64)) # (Batch, SeqLen, Features)
out, c = tf_lstm(inp, dtype=tf.float32)
tf_weight = tf_lstm.weights[0]
tf_bias = tf_lstm.weights[1]

# initialize weights
sess = tf.Session()
init = tf.global_variables_initializer()
sess.run(init)

# tf forward pass
a = np.random.randn(1,50,64).astype(np.float32) # input
b = np.zeros(tf_bias.shape) # set lstm bias to zero
tf_out, lstm_weight, lstm_bias = sess.run([out, tf_weight, tf_bias], {inp: a, tf_bias: b})
assert (lstm_bias == 0).all() # make sure lstm bias was 0

# setup pytorch LSTM
torch_lstm = torch.nn.LSTM(input_size=64, hidden_size=100, num_layers=1, bias=False, batch_first=True)

# set torch weights same as tensorflow weights (is this correct?)
lstm_weight = lstm_weight.T # transpose because pytorch uses left-multiplication instead of right-multiplication
torch_lstm.weight_ih_l0.data = torch.tensor(lstm_weight[:, 0:64])
torch_lstm.weight_hh_l0.data = torch.tensor(lstm_weight[:, 64:164])

# torch forward pass
torch_out, (hn, cn) = torch_lstm(torch.tensor(a))
torch_out = torch_out.detach().numpy() # convert to numpy for compatibility

# compare
assert torch_out.shape == tf_out.shape
print("np.allclose: ", np.allclose(torch_out, tf_out))
print("normalized difference: ", np.linalg.norm(torch_out - tf_out))
</code></pre>
<p>output:</p>
<pre><code class="lang-auto">np.allclose: False
normalized difference: 11.151812
</code></pre>
<p>I am running on cpu with the following dependencies:</p>
<pre><code class="lang-auto">numpy==1.21.5
tensorflow-gpu==1.14.0
torch==1.11.0
</code></pre>
<p>I am running tensorflow v1, cpu version should work, the python wheel is available <a href="https://pypi.org/project/tensorflow/1.14.0/#files" rel="noopener nofollow ugc">here</a> for python&lt;=3.7.</p>
<p>Any help is very appreciated!</p>