Custom layer from keras to pytorch

aaditya: <p>Coming from TensorFlow background, I am trying to convert a snippet of code of the custom layer from Keras to pytorch.</p>
<p>The custom layer in Keras looks like this:</p>
<pre><code class="lang-auto">class Attention_module(tf.keras.layers.Layer):
    def __init__(self, class_num):
        super(Attention_module self).__init__(class_num)
        self.class_num = class_num
        self.Ws = None

    def build(self, input_shape):
        embedding_length = int(input_shape[2])

        self.Ws = self.add_weight(shape=(self.class_num, embedding_length),
                                  initializer=tf.keras.initializers.get('glorot_uniform'), trainable=True)

        super(Attention_module, self).build(input_shape)

    def call(self, inputs):

        sentence_trans = tf.transpose(inputs, [0, 2, 1])
        at = tf.matmul(self.Ws, sentence_trans)
        at = tf.math.tanh(at)
        at = K.exp(at - K.max(at, axis=-1, keepdims=True))
        at = at / K.sum(at, axis=-1, keepdims=True)
        v = K.batch_dot(at, inputs)

        return v
</code></pre>
<p>I want to implement the same in the torch; I have already done the forward pass block but am confused about how to do the embedding and weight initialization the same as the above layer in pytorch?</p>
<pre><code class="lang-auto">class Attention_module(torch.nn.Module):
    def __init__(self, class_num):

        # how to initialize weight with same as above keras layer?

    def forward(self, inputs):
        
        sentence_trans = inputs.permute(0, 2, 1)

        at = torch.mm(self.Ws, sentence_trans)
        at = torch.nn.Tanh(at)
        at = torch.exp(at - torch.max(torch.Tensor(at), dim=-1, keepdims=True).values)
        at = at / torch.sum(at, dim = -1, keepdims=True)
        v = torch.einsum('ijk,ikl-&gt;ijl', at, inputs)

        return v
</code></pre>
<p>Thank you!</p>

tom: <p>There are several aspects:</p>
<ul>
<li>the parameter is created by assigning <code>self.Ws</code> as a  <code>nn.Parameter</code> of the appropriate shape in <code>__init__</code> (after <code>super().__init__()</code>).</li>
<li>the equivalent initialization is <code>torch.nn.init.xavier_uniform_</code>, but you may have to adjust the gain parameter. Check the mean and std (and a histogram plot)</li>
<li>you have some funny stuff in the forward mixing <code>torch.nn</code> things where I would have expected <code>torch.nn.functional</code>, also you should never need to use <code> torch.Tensor</code>,</li>
<li>finally, are you spelling out <code>torch.nn.functional.softmax</code> in the lines with torch.exp and the following? I’d probably write it as softmax if you are. I’m not entirely sure why you would to tanh and then softmax, but hey, what do I know of your problem.</li>
</ul>
<p>I didn’t look at the dimension ordering (the permute and others).</p>
<p>Best regards</p>
<p>Thomas</p>

aaditya: <p>Hi, Thank you for your detailed reply.</p>
<p>I am more confused about <code>self.add_weight</code> function from keras.<br>
is it equivalent to nn.Parameter?<br>
An example would help a lot.</p>

tom: <p>The PyTorch way to track trainable parameters is this, take e.g. the Linear module:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/538c30a713a1ee2a3d654c3e1cdf9cc20b7d8c72/torch/nn/modules/linear.py#L85-L100">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/538c30a713a1ee2a3d654c3e1cdf9cc20b7d8c72/torch/nn/modules/linear.py#L85-L100" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/538c30a713a1ee2a3d654c3e1cdf9cc20b7d8c72/torch/nn/modules/linear.py#L85-L100" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/538c30a713a1ee2a3d654c3e1cdf9cc20b7d8c72/torch/nn/modules/linear.py#L85-L100</a></h4>



  <pre class="onebox">    <code class="lang-py">
      <ol class="start lines" start="85" style="counter-reset: li-counter 84 ;">
          <li>    self.weight = Parameter(torch.empty((out_features, in_features), **factory_kwargs))</li>
          <li>    if bias:</li>
          <li>        self.bias = Parameter(torch.empty(out_features, **factory_kwargs))</li>
          <li>    else:</li>
          <li>        self.register_parameter('bias', None)</li>
          <li>    self.reset_parameters()</li>
          <li>
          </li>
<li>def reset_parameters(self) -&gt; None:</li>
          <li>    # Setting a=sqrt(5) in kaiming_uniform is the same as initializing with</li>
          <li>    # uniform(-1/sqrt(in_features), 1/sqrt(in_features)). For details, see</li>
          <li>    # https://github.com/pytorch/pytorch/issues/57109</li>
          <li>    init.kaiming_uniform_(self.weight, a=math.sqrt(5))</li>
          <li>    if self.bias is not None:</li>
          <li>        fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)</li>
          <li>        bound = 1 / math.sqrt(fan_in) if fan_in &gt; 0 else 0</li>
          <li>        init.uniform_(self.bias, -bound, bound)</li>
      </ol>
    </code>
  </pre>


  </article>

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

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

<p>Empty is an uninitialized tensor, then it is initialized (albeit with funny gain parameters, so you don’t want to use these).</p>

Ansh: <p>Why you left Tensorflow?</p>