Feedforward Sequential Memory Network (FSMN) in PyTorch?

Dee: <p>How to  implement the Feedforward Sequential Memory Network (FSMN) in PyTorch? I’m trying to undersatnd how it works and can not find any example.<br>
I was able to find a code in tensorflow but can not understand it.</p><aside class="onebox pdf" data-onebox-src="https://arxiv.org/pdf/1512.08301.pdf">
  <header class="source">

      <a href="https://arxiv.org/pdf/1512.08301.pdf" target="_blank" rel="noopener nofollow ugc">arxiv.org</a>
  </header>

  <article class="onebox-body">
    <a href="https://arxiv.org/pdf/1512.08301.pdf" target="_blank" rel="noopener nofollow ugc"><span class="pdf-onebox-logo"></span></a>

<h3><a href="https://arxiv.org/pdf/1512.08301.pdf" target="_blank" rel="noopener nofollow ugc">1512.08301.pdf</a></h3>

  <p class="filesize">933.62 KB</p>

  </article>

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

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

<pre><code class="lang-auto">import tensorflow as tf
class FSMN(object):
    def __init__(self, memory_size, input_size, output_size, dtype=tf.float32):
        self._memory_size = memory_size
        self._output_size = output_size
        self._input_size = input_size
        self._dtype = dtype
        self._build_graph()

    def _build_graph(self):
        self._W1 = tf.get_variable("fsmnn_w1", [self._input_size, self._output_size], initializer=tf.truncated_normal_initializer(stddev=5e-2, dtype=self._dtype))
        self._W2 = tf.get_variable("fsmnn_w2", [self._input_size, self._output_size], initializer=tf.truncated_normal_initializer(stddev=5e-2, dtype=self._dtype))
        self._bias = tf.get_variable("fsmnn_bias", [self._output_size], initializer=tf.constant_initializer(0.0, dtype=self._dtype))
        self._memory_weights = tf.get_variable("memory_weights", [self._memory_size], initializer=tf.constant_initializer(1.0, dtype=self._dtype))

    def __call__(self, input_data):
        batch_size = input_data.get_shape()[0].value
        num_steps = input_data.get_shape()[1].value

        memory_matrix = []
        for step in range(num_steps):
            left_num = tf.maximum(0, step + 1 - self._memory_size)
            right_num = num_steps - step - 1
            mem = self._memory_weights[tf.minimum(step, self._memory_size)::-1]
            d_batch = tf.pad(mem, [[left_num, right_num]])
            memory_matrix.append([d_batch])
        memory_matrix = tf.concat(0, memory_matrix)

        h_hatt = tf.batch_matmul([memory_matrix] * batch_size, input_data)
        h = tf.batch_matmul(input_data, [self._W1] * batch_size)
        h += tf.batch_matmul(h_hatt, [self._W2] * batch_size) + self._bias
        return h
</code></pre>

Dee: <p>I converted the code above to Pytorch.</p>
<pre><code class="lang-auto">class FSMN(nn.Module):
    def __init__(self, memory_size, input_size, output_size, dtype=torch.float32):
        super().__init__()
        self._memory_size = memory_size
        self._output_size = output_size
        self._input_size = input_size
        self._dtype = dtype
        self._build_graph()

    def _build_graph(self): 
        self._W1 = Parameter(torch.randn(self._input_size, self._output_size))
        self._W2 = Parameter(torch.randn(self._input_size, self._output_size))
        self._bias = Parameter(torch.randn( self._output_size))  
        self._memory_weights = torch.arange(self._memory_size, dtype=torch.float) 

    def forward(self, input_data):
        batch_size = input_data.size(0) 
        num_steps = input_data.size(1) 

        memory_matrix = []
        for step in range(num_steps):
            left_num = max(0, step + 1 - self._memory_size)
            right_num = num_steps - step - 1
            mem = self._memory_weights[0:min(step, self._memory_size)+1].flip(-1)
            d_batch = F.pad(mem, (left_num, right_num) )
            memory_matrix.append(d_batch)

        memory_matrix = torch.stack((memory_matrix), 0)
        h_hatt = torch.matmul(memory_matrix, input_data)
        h = torch.matmul(input_data, self._W1 )
        h += torch.matmul(h_hatt, self._W2) + self._bias
        return h

model = FSMN(4,8,12) # memory_size x input_size x output_size
input=torch.randn((2,10, 8)) # batch x seq x input 
output = rnn(input)
#output.shape 2 x 10 x 12
</code></pre>

Dee: <aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/d5555/FSMN">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/d5555/FSMN" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/345;"><img src="https://opengraph.githubassets.com/992eb2d1f48c6a39ffd989e2c3a4b96b68a91aa392c0857addf1a5bcc2fea959/d5555/FSMN" class="thumbnail" width="690" height="345"></div>

<h3><a href="https://github.com/d5555/FSMN" target="_blank" rel="noopener nofollow ugc">GitHub - d5555/FSMN: pytorch FSMN</a></h3>

  <p>pytorch FSMN. Contribute to d5555/FSMN development by creating an account on GitHub.</p>


  </article>

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

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