How does Pytorch pack LSTM weights to pass to cuDNN?

0xFFFFFFFF: <p>I have trained a model that uses <code>nn.LSTMCell</code>. For reasons of throughput, I want to directly call cuDNN’s <code>cudnnRNNForwardInference</code>. So I have to export weights of <code>nn.LSTMCell</code>. For other layers such as linear or convolution, this wasn’t hard. However, this is very difficult for <code>nn.LSTMCell</code> because <code>nn.LSTMCell</code> takes 2 sets of weights and 2 sets of biases, while <code>cudnnRNNForwardInference</code> only takes a single set of weight and a single set of bias:</p>
<h3><code>nn.LSTMCell</code></h3>
<pre><code class="lang-auto">~LSTMCell.weight_ih – the learnable input-hidden weights, of shape (4*hidden_size, input_size)

~LSTMCell.weight_hh – the learnable hidden-hidden weights, of shape (4*hidden_size, hidden_size)

~LSTMCell.bias_ih – the learnable input-hidden bias, of shape (4*hidden_size)

~LSTMCell.bias_hh – the learnable hidden-hidden bias, of shape (4*hidden_size)
</code></pre>
<h3><code>cudnnRNNForwardInference</code></h3>
<pre><code class="lang-auto">cudnnStatus_t cudnnRNNForwardInference(
    cudnnHandle_t                   handle,
    const cudnnRNNDescriptor_t      rnnDesc,
    const int                       seqLength,
    const cudnnTensorDescriptor_t  *xDesc,
    const void                     *x,
    const cudnnTensorDescriptor_t   hxDesc,
    const void                     *hx,
    const cudnnTensorDescriptor_t   cxDesc,
    const void                     *cx,
    const cudnnFilterDescriptor_t   wDesc,
    const void                     *w,
    const cudnnTensorDescriptor_t   *yDesc,
    void                           *y,
    const cudnnTensorDescriptor_t   hyDesc,
    void                           *hy,
    const cudnnTensorDescriptor_t   cyDesc,
    void                           *cy,
    void                           *workspace,
    size_t                          workSpaceSizeInBytes)
</code></pre>
<p>How would I pack each of the weights: <code>weight_ih</code>, <code>weight_hh</code>,  <code>bias_ih</code>, <code>bias_hh</code> s.t. it can be passed to cuDNN’s <code>cudnnRNNForwardInference</code> . as <code>wDesc</code> parameter?</p>

ZoroDerVonCodier: <p>See <a href="https://kaixih.github.io/keras-cudnn-rnn/" rel="noopener nofollow ugc">Parameters In Tensorflow Keras RNN and CUDNN RNN - Kaixi Hou’s Log</a> - great write up on this!</p>

ZoroDerVonCodier: <p>this function will pack them in the order used by cudnn as per the referenced article above and the pytorch documents at <a href="https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html" rel="noopener nofollow ugc">LSTM — PyTorch 2.0 documentation</a>.  Note pytorch switches the ordering of the wtig and wtio from the previous article which uses karas.</p>
<h1>
<a name="convert-lstm-weights-to-cudnn-format-1" class="anchor" href="#convert-lstm-weights-to-cudnn-format-1"></a>convert lstm weights to cudnn format</h1>
<pre><code>def get_cudnn_lstm_weights(lstm):
    num_layers = lstm.num_layers
    hidden_size = lstm.hidden_size
    wts = []
    for param in lstm.state_dict():
        param_val = lstm.state_dict()[param]
        wts.append(param_val)

    all_wts = []
    idx = 0
    for i in range(0, num_layers):
        wtii = wts[idx][:hidden_size,:]
        wtif = wts[idx][hidden_size:hidden_size*2,:]
        wtig = wts[idx][hidden_size*2:hidden_size*3,:]
        wtio = wts[idx][hidden_size*3:hidden_size*4,:]
        idx = idx + 1

        wthi = wts[idx][:hidden_size,:]
        wthf = wts[idx][hidden_size:hidden_size*2,:]
        wthg = wts[idx][hidden_size*2:hidden_size*3,:]
        wtho = wts[idx][hidden_size*3:hidden_size*4,:]
        idx = idx + 1

        bii = wts[idx][:hidden_size]
        bif = wts[idx][hidden_size:hidden_size*2]
        big = wts[idx][hidden_size*2:hidden_size*3]
        bio = wts[idx][hidden_size*3:hidden_size*4]
        idx = idx + 1

        bhi = wts[idx][:hidden_size]
        bhf = wts[idx][hidden_size:hidden_size*2]
        bhg = wts[idx][hidden_size*2:hidden_size*3]
        bho = wts[idx][hidden_size*3:hidden_size*4]
        idx = idx + 1

        wts1 = [wtii, wtif, wtio, wtig, wthi, wthf, wtho, wthg]
        b1 = [bii, bif, bio, big, bhi, bhf, bho, bhg]

        shape = [-1]
        weights = [torch.reshape(torch.transpose(x, 0, 1), shape) for x in wts1]
        biases = [torch.reshape(x, shape) for x in b1]
        cudnnwts = torch.concat(weights + biases, axis=0)
        all_wts.append(cudnnwts)

    full_cudnnwts = torch.concat(all_wts, axis=0) if num_layers &gt; 1 else all_wts[0]
    return full_cudnnwts
</code></pre>