JIT GRU Implementation Optimization

Maghoumi: <p>Following <a href="https://pytorch.org/blog/optimizing-cuda-rnn-with-torchscript/" rel="nofollow noopener">this guide</a> and <a href="https://github.com/pytorch/pytorch/blob/master/benchmarks/fastrnns/custom_lstms.py" rel="nofollow noopener">this code</a> and inspired by <a href="https://github.com/pytorch/pytorch/issues/5261" rel="nofollow noopener">this discussion</a>, I whipped up my own GRU cell implementation using JIT as follows.</p>
<pre><code class="lang-auto">class JitGRUCell(jit.ScriptModule):
    def __init__(self, input_size, hidden_size):
        super(JitGRUCell, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.weight_ih = Parameter(torch.Tensor(3 * hidden_size, input_size))
        self.weight_hh = Parameter(torch.Tensor(3 * hidden_size, hidden_size))
        self.bias_ih = Parameter(torch.Tensor(3 * hidden_size))
        self.bias_hh = Parameter(torch.Tensor(3 * hidden_size))

        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1.0 / math.sqrt(self.hidden_size)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, stdv)

    @jit.script_method
    def forward(self, x, hidden):
        # type: (Tensor, Tensor) -&gt; Tensor
        x = x.view(-1, x.size(1))
        x_results = torch.mm(x, self.weight_ih.t()) + self.bias_ih
        h_results = torch.mm(hidden, self.weight_hh.t()) + self.bias_hh

        x_results = x_results.squeeze()
        h_results = h_results.squeeze()

        i_r, i_z, i_n = x_results.chunk(3, 1)
        h_r, h_z, h_n = h_results.chunk(3, 1)

        r = torch.sigmoid(i_r + h_r)
        z = torch.sigmoid(i_z + h_z)
        n = torch.tanh(i_n + r * h_n)

        return n - torch.mul(n, z) + torch.mul(z, hidden)
</code></pre>
<p>The implementation itself looks very straightforward, however performance is not great.</p>
<p>What are some tricks that I could use to optimize this implementation further? I’m very new to JIT script modules and I appreciate each and every suggestion.</p>

wuclddwf: <p>hi, Maghoumi. I’ve suffered form the same problem as you， do you fix it?</p>

Maghoumi: <p>Sorry for the super late response. I just randomly saw this thread again. Nope… Still using the same code. Did you solve this problem?</p>

Maghoumi: <p>What do you mean “return h”? IIRC the output of a GRU cell is the hidden state itself (you feed it back to itself as the previous hidden state)</p>
<p>Also, I had some unit tests in place, and could confirm that this code returned close-enough results to PyTorch’s GRUCell implementation</p>

Dan_Erez: <p>yep. you are right. good thing i started with i moght be wrong …<img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

Maghoumi: <p>By the way, here’s my full implementation (including support for multiple layers, etc.): <a href="https://github.com/Maghoumi/JitGRU" rel="nofollow noopener">https://github.com/Maghoumi/JitGRU</a></p>
<p>Suggestions for improvement are greatly welcome!</p>

G.M: <p>If you are comparing jit version to default version, there certainly will be a gap, especially during backward process. I noticed that both the implementation from <a href="https://github.com/pytorch/pytorch/blob/master/benchmarks/fastrnns/custom_lstms.py" rel="nofollow noopener">here</a> and your implementation used <code>+=[i]</code>, but I recommend using <code>.append(i)</code>. I don’t know how are these two methods translated to C++ but it seems that <code>.append(i)</code> is faster than simply adding, especially when you want to add many elements.</p>