How does Neural Network as Input work?

Vishu_Gupta: <p>Recently I have been studying GNN and I came across the following piece of code which I am unable to understand:</p>
<pre><code class="lang-auto">self.pooling = nn.ModuleList([WeightedAttention(
            gate_nn=SimpleNetwork(2*fea_len, 1, hidden_ele),
            message_nn=SimpleNetwork(2*fea_len, fea_len, hidden_msg),
            # message_nn=nn.Linear(2*fea_len, fea_len),
            # message_nn=nn.Identity(),
            ) for _ in range(num_heads)])
</code></pre>
<p>Is the <code>WeightedAttention</code> taking <code>Simple Network</code> as Input which comprises of linear layer? Also, how does it work inside the <code>WeightedAttention</code> class? The <code>WeightedAttention</code> and <code>Simple Network</code> class are as follows:</p>
<pre><code class="lang-auto">class WeightedAttention(nn.Module):
    """
    Weighted softmax attention layer
    """
    def __init__(self, gate_nn, message_nn, num_heads=1):
        """
        Inputs
        ----------
        gate_nn: Variable(nn.Module)
        """
        
        super(WeightedAttention, self).__init__()
        self.gate_nn = gate_nn
        self.message_nn = message_nn
        self.pow = torch.nn.Parameter(torch.randn((1)))

    def forward(self, fea, index, weights):
        """ forward pass """

        gate = self.gate_nn(fea)

        gate = gate - scatter_max(gate, index, dim=0)[0][index]
        gate = (weights ** self.pow) * gate.exp()
        # gate = weights * gate.exp()
        # gate = gate.exp()
        gate = gate / (scatter_add(gate, index, dim=0)[index] + 1e-13)

        fea = self.message_nn(fea)
        out = scatter_add(gate * fea, index, dim=0)

        return out

    def __repr__(self):
        return '{}(gate_nn={})'.format(self.__class__.__name__,
                                       self.gate_nn)


class SimpleNetwork(nn.Module):
    """
    Simple Feed Forward Neural Network
    """
    def __init__(self, input_dim, output_dim, hidden_layer_dims):
        """
        Inputs
        ----------
        input_dim: int
        output_dim: int
        hidden_layer_dims: list(int)

        """
        super(SimpleNetwork, self).__init__()

        dims = [input_dim]+hidden_layer_dims
        
        self.fcs = nn.ModuleList([nn.Linear(dims[i], dims[i+1])
                                  for i in range(len(dims)-1)])
        self.acts = nn.ModuleList([nn.LeakyReLU() for _ in range(len(dims)-1)])

        self.fc_out = nn.Linear(dims[-1], output_dim)

    def forward(self, fea):
        for fc, act in zip(self.fcs, self.acts):
            fea = act(fc(fea))

        return self.fc_out(fea)

    def __repr__(self):
        return '{}'.format(self.__class__.__name__)
</code></pre>

ptrblck: <p><code>WeightedAttention</code> uses a <code>SimpleNetwork</code> for <code>gate_nn</code> and another one for <code>message_nn</code>.<br>
The usage of these modules is shown in <code>WeightedAttention.forward</code>.<br>
I don’t know, how <code>scatter_max</code> and <code>scatter_add</code> are implemented, but the usage of the <code>SimpleNetwork</code> modules is:</p>
<pre><code class="lang-python">gate = self.gate_nn(fea)
...
fea = self.message_nn(fea)
...
</code></pre>
<p><code>SimpleNetwork</code> uses an <code>nn.ModulesList</code> with several linear layers, another module list to store the activation functions, as well as a final linear layer.</p>

Vishu_Gupta: <p>But how can you perform <code>scatter_max</code> or <code>scatter_add</code> to a neural network? What will it affect, the weight of the neural network or something else?</p>

ptrblck: <p><code>scatter_add</code> can be understood as an advanced indexing method.<br>
Have a look at the <a href="https://pytorch.org/docs/stable/tensors.html?highlight=scatter#torch.Tensor.scatter_add_">docs</a> with the example:</p>
<pre><code class="lang-python">&gt;&gt;&gt; x = torch.rand(2, 5)
&gt;&gt;&gt; x
tensor([[ 0.3992,  0.2908,  0.9044,  0.4850,  0.6004],
        [ 0.5735,  0.9006,  0.6797,  0.4152,  0.1732]])
&gt;&gt;&gt; torch.zeros(3, 5).scatter_(0, torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]]), x)
tensor([[ 0.3992,  0.9006,  0.6797,  0.4850,  0.6004],
        [ 0.0000,  0.2908,  0.0000,  0.4152,  0.0000],
        [ 0.5735,  0.0000,  0.9044,  0.0000,  0.1732]])

&gt;&gt;&gt; z = torch.zeros(2, 4).scatter_(1, torch.tensor([[2], [3]]), 1.23)
&gt;&gt;&gt; z
tensor([[ 0.0000,  0.0000,  1.2300,  0.0000],
        [ 0.0000,  0.0000,  0.0000,  1.2300]])
</code></pre>
<p><code>scatter_max</code> isn’t a PyTorch method, if I’m not mistaken, but I assume it’ll use the <code>max</code> operation instead of adding the values.</p>
<p>I would recommend to print all tensors, which are used in these operations, and to check the output of the operations in question.</p>
<p>These operations won’t affect the weights or other parametes directly, but are most likely used for some indexing.</p>