Batch size handling with LSTM

Ahmed_Akl · August 11, 2021, 2:07pm

Hi All,
I am trying to implement custom LSTM layer with custom cell.
It is working OK when I pass only one sample, but when I want to pass a batch of data a problem appear.
For example:
If my data look like [128, 64] where 128 is max sequence length and 64 is word embedding size, the model works fine
I want to know how to handle a batch of data e.g. [16, 128, 64] ?
I search for the implementation of the lstm cell in pytorch , but I can’t find the code that handles the batch

github.com

pytorch/pytorch/blob/f446e835ee35fed7e2f656827f3708a5044c4fa3/torch/nn/modules/rnn.py#L689

    
      
                                    max_batch_size, self.hidden_size,
                                    dtype=input.dtype, device=input.device)
              hx = (h_zeros, c_zeros)
          else:
              # Each batch of the hidden state should match the input sequence that
              # the user believes he/she is passing in.
              hx = self.permute_hidden(hx, sorted_indices)
          
          
self.check_forward_args(input, hx, batch_sizes)
          if batch_sizes is None:
              result = _VF.lstm(input, hx, self._flat_weights, self.bias, self.num_layers,
                                self.dropout, self.training, self.bidirectional, self.batch_first)
          else:
              result = _VF.lstm(input, batch_sizes, hx, self._flat_weights, self.bias,
                                self.num_layers, self.dropout, self.training, self.bidirectional)
          output = result[0]
          hidden = result[1:]
          # xxx: isinstance check needs to be in conditional for TorchScript to compile
          if isinstance(orig_input, PackedSequence):
              output_packed = PackedSequence(output, batch_sizes, sorted_indices, unsorted_indices)
              return output_packed, self.permute_hidden(hidden, unsorted_indices)

and this code leads to this impl.

github.com

pytorch/pytorch/blob/f446e835ee35fed7e2f656827f3708a5044c4fa3/test/test_jit.py#L2640

    
      
                  return x + 10
          
          
    with TemporaryFileName() as fname:
                  fn.save(fname)
                  with io.open(fname, 'rb') as f:
                      self.assertTrue(torch.serialization._is_zipfile(f))
          
          
def test_python_bindings(self):
              lstm_cell = torch.jit.script(LSTMCellS)
          
          
    def lstm(x, hx, cx, w_ih, w_hh, b_ih, b_hh):
                  for i in range(x.size(0)):
                      hx, cx = lstm_cell(x[i], hx, cx, w_ih, w_hh, b_ih, b_hh)
                  return hx
          
          
    slstm = torch.jit.script(lstm)
          
          
    inputs = get_lstm_inputs('cpu', training=True, seq_length=10)
              slstm(*inputs).sum().backward()
              global fw_graph
              fw_graph = slstm.graph_for(*inputs)

is pytorch for loop the data batch?
Thanks in advance.

Ahmed_Akl · September 1, 2021, 1:13pm

Any help will be appreciated @ptrblck

AbdulsalamBande · October 15, 2021, 9:51pm

@Ahmed_Akl , 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 torch.nn.RNN Module explained - YouTube