Batch size handling with LSTM

Ahmed_Akl: <p>Hi All,<br>
I am trying to implement custom LSTM layer with custom cell.<br>
It is working OK when I pass only one sample, but when I want to pass a batch of data a problem appear.<br>
For example:<br>
If my data look like [128, 64] where 128 is max sequence length and 64 is word embedding size, the model works fine<br>
I want to know how to handle a batch of data e.g. [16, 128, 64] ?<br>
I search for the implementation of the lstm cell in pytorch , but I can’t find the code that handles the batch</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/f446e835ee35fed7e2f656827f3708a5044c4fa3/torch/nn/modules/rnn.py#L689">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/f446e835ee35fed7e2f656827f3708a5044c4fa3/torch/nn/modules/rnn.py#L689" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/f446e835ee35fed7e2f656827f3708a5044c4fa3/torch/nn/modules/rnn.py#L689" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/f446e835ee35fed7e2f656827f3708a5044c4fa3/torch/nn/modules/rnn.py#L689</a></h4>



  <pre class="onebox">    <code class="lang-py">
      <ol class="start lines" start="679" style="counter-reset: li-counter 678 ;">
          <li>                          max_batch_size, self.hidden_size,</li>
          <li>                          dtype=input.dtype, device=input.device)</li>
          <li>    hx = (h_zeros, c_zeros)</li>
          <li>else:</li>
          <li>    # Each batch of the hidden state should match the input sequence that</li>
          <li>    # the user believes he/she is passing in.</li>
          <li>    hx = self.permute_hidden(hx, sorted_indices)</li>
          <li>
          </li>
<li>self.check_forward_args(input, hx, batch_sizes)</li>
          <li>if batch_sizes is None:</li>
          <li class="selected">    result = _VF.lstm(input, hx, self._flat_weights, self.bias, self.num_layers,</li>
          <li>                      self.dropout, self.training, self.bidirectional, self.batch_first)</li>
          <li>else:</li>
          <li>    result = _VF.lstm(input, batch_sizes, hx, self._flat_weights, self.bias,</li>
          <li>                      self.num_layers, self.dropout, self.training, self.bidirectional)</li>
          <li>output = result[0]</li>
          <li>hidden = result[1:]</li>
          <li># xxx: isinstance check needs to be in conditional for TorchScript to compile</li>
          <li>if isinstance(orig_input, PackedSequence):</li>
          <li>    output_packed = PackedSequence(output, batch_sizes, sorted_indices, unsorted_indices)</li>
          <li>    return output_packed, self.permute_hidden(hidden, unsorted_indices)</li>
      </ol>
    </code>
  </pre>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>and this code leads to this impl.</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/f446e835ee35fed7e2f656827f3708a5044c4fa3/test/test_jit.py#L2640">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/f446e835ee35fed7e2f656827f3708a5044c4fa3/test/test_jit.py#L2640" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/f446e835ee35fed7e2f656827f3708a5044c4fa3/test/test_jit.py#L2640" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/f446e835ee35fed7e2f656827f3708a5044c4fa3/test/test_jit.py#L2640</a></h4>



  <pre class="onebox">    <code class="lang-py">
      <ol class="start lines" start="2630" style="counter-reset: li-counter 2629 ;">
          <li>        return x + 10</li>
          <li>
          </li>
<li>    with TemporaryFileName() as fname:</li>
          <li>        fn.save(fname)</li>
          <li>        with io.open(fname, 'rb') as f:</li>
          <li>            self.assertTrue(torch.serialization._is_zipfile(f))</li>
          <li>
          </li>
<li>def test_python_bindings(self):</li>
          <li>    lstm_cell = torch.jit.script(LSTMCellS)</li>
          <li>
          </li>
<li class="selected">    def lstm(x, hx, cx, w_ih, w_hh, b_ih, b_hh):</li>
          <li>        for i in range(x.size(0)):</li>
          <li>            hx, cx = lstm_cell(x[i], hx, cx, w_ih, w_hh, b_ih, b_hh)</li>
          <li>        return hx</li>
          <li>
          </li>
<li>    slstm = torch.jit.script(lstm)</li>
          <li>
          </li>
<li>    inputs = get_lstm_inputs('cpu', training=True, seq_length=10)</li>
          <li>    slstm(*inputs).sum().backward()</li>
          <li>    global fw_graph</li>
          <li>    fw_graph = slstm.graph_for(*inputs)</li>
      </ol>
    </code>
  </pre>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>is pytorch for loop the data batch?<br>
Thanks in advance.</p>

Ahmed_Akl: <p>Any help will be appreciated <a class="mention" href="/u/ptrblck">@ptrblck</a></p>

AbdulsalamBande: <p><a class="mention" href="/u/ahmed_akl">@Ahmed_Akl</a> , You can do it ,but keep in my that the states such as the hidden and the cell states rely on the batch size. Understanding this video might help answer your question <a href="https://www.youtube.com/watch?v=4p9SXDSwxqA&amp;ab_channel=PytorchModulesExplained" class="inline-onebox" rel="noopener nofollow ugc">torch.nn.RNN Module explained - YouTube</a></p>