Rewriting the code of Tensorflow (Need Review) - Just 2 lines only

krrishabh: <p>In TensorFlow, there is two function get_initial_state(self) and single_step(self, context, word_id)</p>
<p>where (Link: <a href="https://github.com/kaldi-asr/kaldi/blob/2c7e78f0757120a965e671f2843c7c07ac7141d4/egs/wsj/s5/steps/tfrnnlm/lstm.py#L115-L139" rel="nofollow noopener">https://github.com/kaldi-asr/kaldi/blob/2c7e78f0757120a965e671f2843c7c07ac7141d4/egs/wsj/s5/steps/tfrnnlm/lstm.py#L115-L139</a> )</p>
<p><span class="mention">@tf.function</span><br>
def get_initial_state(self):<br>
“”“Exported function which emits zeroed RNN context vector.”""<br>
# This seems a bug in TensorFlow, but passing tf.int32 makes the state tensor also int32.<br>
fake_input = tf.constant(0, dtype=tf.float32, shape=[1, 1])<br>
initial_state = tf.stack(self.rnn.get_initial_state(fake_input))<br>
return {“initial_state”: initial_state}</p>
<p><span class="mention">@tf.function</span><br>
def single_step(self, context, word_id):<br>
“”“Exported function which perform one step of the RNN model.”""<br>
rnn = tf.keras.layers.RNN(self.cells, return_state=True)<br>
context = tf.unstack(context)<br>
context = [tf.unstack© for c in context]</p>
<pre><code>inputs = self.embedding(word_id)
rnn_out_and_states = rnn(inputs, initial_state=context)

rnn_out = rnn_out_and_states[0]
rnn_states = tf.stack(rnn_out_and_states[1:])

logits = self.fc(rnn_out)
output = self.get_score(logits)
log_prob = output[0, word_id[0, 0]]
return {"log_prob": log_prob, "rnn_states": rnn_states, "rnn_out": rnn_out}
</code></pre>
<p>where these functions are used in Kaldi Language Model rescoring (which is present in the same code from line 230-236)</p>
<h1>Export</h1>
<p>print(“Saving model to %s.” % FLAGS.save_path)<br>
spec = [tf.TensorSpec(shape=[config.num_layers, 2, 1, config.hidden_size], dtype=data_type(), name=“context”),<br>
tf.TensorSpec(shape=[1, 1], dtype=tf.int32, name=“word_id”)]<br>
cfunc = model.single_step.get_concrete_function(*spec)<br>
cfunc2 = model.get_initial_state.get_concrete_function()<br>
tf.saved_model.save(model, FLAGS.save_path, signatures={“single_step”: cfunc, “get_initial_state”: cfunc2})</p>
<p>To which I did in PyTorch LM (standard example)<br>
(Link: <a href="https://github.com/pytorch/examples/blob/master/word_language_model/main.py" rel="nofollow noopener">https://github.com/pytorch/examples/blob/master/word_language_model/main.py</a>)<br>
traced_script_module = torch.jit.trace(model, (train_data, hidden))<br>
traced_script_module.save(‘newmodel.pt’)</p>
<p>Does This do the same as Tesorflow both functions do?<br>
Can I get initial_state, rnn_state, and rnn_output?</p>
<p>Thank you in advance (Very Important it is)</p>

ptrblck: <p>Could you explain, what the TF methods are doing?<br>
I would recommend to store fixed numpy inputs, load them for both frameworks (in PyTorch you would have to transform the numpy array to a tensor via <code>torch.from_numpy</code>), and compare their outputs.<br>
This would make sure that both methods are indeed using the same underlying operations and parameters.</p>

krrishabh: <p>Sir the first function get_initial_state() - gives initial_state (<strong>h_0</strong>)<br>
whereas second function single_step(self, context, word_id) - gives log_prob (output probability) rnn_states (<strong>h_(n-1):h_1</strong>) and output (<strong>h_n</strong>)</p>
<p>Sir correct me if I am wrong. Sir I need these in Pytorch…<br>
So, I have stored…<br>
traced_script_module = torch.jit.trace(model, (train_data, hidden))<br>
traced_script_module.save(‘newmodel.pt’)</p>
<p>I will be using it in C++</p>
<p>It is for Pytorch-Kaldi Language Model rescoring…</p>

ptrblck: <p>The mentioned TF functions seem to be related to some internal utility functions for RNNs, while your PyTorch code traces a model and stores it.<br>
I’m not sure, how these approaches are related.</p>
<p>If your PyTorch model also implements the TF methods, then you should be probably fine.</p>

krrishabh: <p>Sir any suggestion to implement TF methods in Pytorch. Sorry Sir to disturb you again chasing a deadline “Interspeech”</p>
<p>Moreover Sir, We need to call all in C++. (Model + Methods Result) because Kaldi part except only C++ code for lattice rescoring.</p>

ptrblck: <p>As I’m not deeply familiar with the TF code, I assume <code>get_initial_state</code> might be equivalent to <a href="https://github.com/pytorch/examples/blob/b9f3b2ebb9464959bdbf0c3ac77124a704954828/word_language_model/model.py#L56-L62">init_hidden</a> and <code>single_step</code> seems to refer to the <a href="https://github.com/pytorch/examples/blob/b9f3b2ebb9464959bdbf0c3ac77124a704954828/word_language_model/model.py#L48-L54">forward method</a>.</p>
<p>If you need exactly the same results, I would still recommend to run the methods in isolation with a fixed input array and make sure you get the desired outputs.</p>

krrishabh: <p>Sir It is <code>get_initial_state</code>  is equivalent to <a href="https://github.com/pytorch/examples/blob/b9f3b2ebb9464959bdbf0c3ac77124a704954828/word_language_model/model.py#L56-L62" rel="nofollow noopener">init_hidden</a> and  <code>single_step</code>  is <a href="https://github.com/pytorch/examples/blob/b9f3b2ebb9464959bdbf0c3ac77124a704954828/word_language_model/model.py#L48-L54" rel="nofollow noopener">forward method</a></p>

krrishabh: <p>That’s why I used…</p>
<p><code>traced_script_module = torch.jit.trace(model, (train_data, hidden))</code><br>
<code>traced_script_module.save(‘newmodel.pt’)</code></p>
<p>And though to use <strong>forward</strong> function in C++.</p>

krrishabh: <p>I had to write these functions because it was used in a Wrapper for extracting getting the initial state.</p>
<p>I have written that: <a href="https://github.com/cyfer0618/kaldi-pytorch-rnnlm/blob/cbaffa572adf2a8006315ffb84fdcbf0318d5670/kaldi/egs/MemoryLM/Gujarati-Recipe/steps/pyrnnlm/model.py#L65-L72" rel="nofollow noopener">Link</a></p>
<pre><code class="lang-auto">    @torch.jit.export
    def get_initial_state(self, bsz):
        #bsz -&gt; batch_size
        #h,c = self.init_hidden(bsz)
        batchsize=bsz.item()
        h=torch.zeros(self.nlayers, batchsize, self.nhid)
        c=torch.zeros(self.nlayers, batchsize, self.nhid)
        return h
</code></pre>