Hi,
I’m currently testing a variant of the LSTM architecture called subLSTM. I was trying to get an efficient implementation to speed up my tests since my PyTorch implemenation its still very slow compared to the library LSTM. I also tried using TorchScript but its still much slower than the LSTM version. Specifically I used jit.script to compile the inner loop of the RNN (similar to the LSTM implementation). So my questions are:

Why does this not make it work as fast? Is it a problem with the way I’m using jit.script?
I was trying to access the RNN implementations in _C._VariableFunctions, but I can’t see it. Is there any way of modifying that code? The modifications I need are trivial, and while I could implement it myself I may not get such an efficient code since I am not that experienced with CUDA (and my C++ is pretty rusty).

Here is a link to my code:

  
      github.com
  
  
    mllera14/subLSTM/blob/optimize/src/subLSTM/nn.py 53
import math
from itertools import product
import numbers

import torch
import torch.nn as nn
import torch.jit as jit
from torch.nn.modules.activation import Sigmoid
# from torch.nn.modules.rnn import PackedSequence
# from torch.nn.utils.rnn import pack_padded_sequence as pack, pad_packed_sequence as pad

from .functional import sublstm, fsublstm

@jit.script
def _sublstm_forward(input, hx, weights, num_layers):
    # type: (List[Tensor], List[Tuple[Tensor, Tensor]], List[List[Tensor]], int) -> Tuple[Tensor, List[Tuple[Tensor, Tensor]]]
    timesteps = len(input)

    for time in range(timesteps):
        for l in range(num_layers):


  This file has been truncated. show original

  
  
    
    
  
  



  
      github.com
  
  
    mllera14/subLSTM/blob/optimize/src/subLSTM/functional.py 27
import torch
import torch.nn as nn
import torch.nn.functional as F

@torch.jit.script
def sublstm(input, hidden, W, R, bi, bh):
    # type: (Tensor, Tuple[Tensor, Tensor], Tensor, Tensor, Tensor, Tensor) -> Tuple[Tensor, Tensor]
    h_tm1, c_tm1 = hidden

    proj_input = torch.sigmoid(F.linear(input, W , bi) + F.linear(h_tm1, R, bh))
    in_gate, out_gate, z_t, f_gate = proj_input.chunk(4, 1)

    c_t = c_tm1 * f_gate + z_t - in_gate
    h_t = torch.sigmoid(c_t) - out_gate

    return h_t, c_t

@torch.jit.script
def fsublstm(input, hidden, W, R, bi, bh, f_gate):
    # type: (Tensor, Tuple[Tensor, Tensor], Tensor, Tensor, Tensor, Tensor, Tensor) -> Tuple[Tensor, Tensor]


  This file has been truncated. show original

  
  
    
    
  
  


Thanks,