How do I process a batch in my forward() function?

agt: <p>I’m making a module and I expected to get 1 input (shape (2,2,3,3)) at a time. I just realized <code>nn.Linear</code> expects a batch dimension, so I need to expect batches not individual inputs. How does <code>nn.Linear</code> process batches and how can I process batches in my forward()? Would I just put everything in a loop over the batch elements?</p>
<p>This is what my forward() method looks like:</p>
<pre><code class="lang-auto">    def forward(self, pair_of_graphs):
        embeddings = []

        for graph in pair_of_graphs:
            node_matrix, adjacency_matrix = graph
            steps = 5
            for step in range(steps):
                message_passed_node_matrix = torch.matmul(adjacency_matrix, node_matrix)
                alpha = torch.nn.Parameter(torch.zeros(1))
                node_matrix = alpha*node_matrix + (1-alpha)*messaged_passed_node_matrix

                new_node_matrix = torch.zeros(len(node_matrix), self.linear_2.in_features)
                for node_i in range(len(node_matrix)):
                    linear_layer = self.linear_1 if step == 0 else self.linear_2
                    new_node_matrix[node_i] = linear_layer(node_matrix[node_i])
                node_matrix = new_node_matrix
                
            weights_for_average = torch.zeros(len(node_matrix))
            for node_i in range(len(node_matrix)):
                weights_for_average[node_i] = self.linear_3(node_matrix[node_i])

            weighted_sum_node = torch.matmul(weights_for_average, node_matrix)
            embeddings.append(weighted_sum_node)

        concat = torch.cat(embeddings)
        out = self.linear_4(concat)

        return out
</code></pre>
<p>Any other suggestions about how I am handling my input (e.g. looping through each node of each graph) would be appreciated too.</p>

ptrblck: <p>All PyTorch layers accept and expect batched inputs and don’t need a for loop or any other change.<br>
<code>nn.Linear</code> in particular expects the input to have the shape <code>[batch_size, *, in_features]</code>, where the <code>*</code> is a variable number of dimensions. The linear layer will be applied to each sample in these variable number of dimensions as if you would be using a loop over these dims.</p>

agt: <p>But I’m not only passing my input into pytorch layers. I also have a customer parameter <code>alpha</code> that I use to combine my input with an output of a linear layer: <code>alpha*input + (1-alpha)*output</code>. I have other custom behavior like that too.</p>
<p>In that case, would I have to manually loop over the input batch and build the new tensor after doing my custom computations on each item of the batch? Or is there a better way to do these sorts of things?</p>

ptrblck: <aside class="quote no-group" data-username="agt" data-post="3" data-topic="100304">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/7feea3/48.png" class="avatar"> agt:</div>
<blockquote>
<p>But I’m not only passing my input into pytorch layers.</p>
</blockquote>
</aside>
<p>It seems you still passing <code>input</code> to the layer and apply the <code>alpha</code> parameter later or could you explain how you are passing <code>alpha</code> to e.g. the linear layer?</p>
<aside class="quote no-group" data-username="agt" data-post="3" data-topic="100304">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/7feea3/48.png" class="avatar"> agt:</div>
<blockquote>
<p>In that case, would I have to manually loop over the input batch and build the new tensor after doing my custom computations on each item of the batch?</p>
</blockquote>
</aside>
<p>PyTorch applies broadcasting, so if <code>alpha</code> is a scalar tensor you could directly run the posted line of code.<br>
On the other hand, even if <code>alpha</code> has the shape <code>[batch_size]</code> it should still work (and you might need to <code>unsqueeze()</code> dimensions to enable broadcasting, but it depends on the shapes of the other tensors).</p>