Different model branches for different samples in batch

flpgrz: <p>Hello, imagine I have a NN which has multiple output branches. I want to train the upstream layers jointly, while the output layers can be plugged in, based on a categorical variable. Let’s assume there are simply 2 FC output layers:</p>
<pre><code class="lang-auto">class MyModel(nn.Model):
      def __init__(self):
              ... # some layers 
              output_fc1 = nn.Linear(16,1)
              output_fc2 = nn.Linear(16,1)
</code></pre>
<p>Now, I have a batch of samples <code>x</code>, and an associated tensor <code>output_branch_selector</code>. The elements of <code>output_branch_selector</code> tell me which output layer should be plugged in for each sample in a batch, and they can either be 0 (for <code>output_fc1</code>) or 1 (for <code>output_fc2</code>).</p>
<p>Is there a way to do this in a single forward pass by assigning automatically an output branch based on <code>output_branch_selector</code>? E.g.</p>
<pre><code class="lang-auto">model = MyModel()
output = model(x, output_branch_selector)
</code></pre>

ID56: <p>Let me try to generalize it a bit for you, to a scalable number of branches.</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    "A model with variable number of branches."
    def __init__(self, in_features: int = 16, out_features: int = 1, num_branches: int = 5) -&gt; None:
        """Init method.

        Args:
            in_features (int, optional): Number of input features. Defaults to 16.
            out_features (int, optional): Output features. Defaults to 1.
            num_branches (int, optional): Number of branches in model. Defaults to 5.
        """
        super().__init__()
        self.output_fcs = nn.ModuleList(
            [nn.Linear(in_features, out_features) for _ in range(num_branches)]
        )
        self.num_branches = num_branches
        self.in_features = in_features
        self.out_features = out_features

    def forward(self, X: torch.Tensor, S: torch.Tensor) -&gt; torch.Tensor:
        """Forward Method.

        Args:
            X (torch.Tensor): Input of shape (batch, in_features).
            S (torch.Tensor): Input of shape (batch,). Values in the range (0, num_branches - 1).

        Returns:
            torch.Tensor: Model predictions, of shape (batch, out_features).
        """
        return torch.cat([self.output_fcs[s](X[i, :]).unsqueeze(0) for i, s in enumerate(S)], axis=0)
</code></pre>
<p>Let’s do a simple test for this as well:</p>
<pre><code class="lang-python">model = MyModel(16, 1, 5)
batch_size = 4
dummy_input = torch.randn(batch_size, 16)
dummy_selector = torch.randint(0, 5, (batch_size,))

print(dummy_selector)
# tensor([0, 3, 3, 4])

output = model(dummy_input, dummy_selector)
print(output, output.shape)
# tensor([[-0.8775],
#         [ 0.0129],
#         [-0.3767],
#         [-0.4781]], grad_fn=&lt;CatBackward0&gt;) torch.Size([4, 1])
</code></pre>
<p>So in the above example, no sample should have passed through branch 1 and 2.</p>
<p>Note the forward pass is kind of inefficient now: we lost any advantage of batching process as we’re doing B Linears if we have B samples.</p>
<p>A more efficient alternative would be to calculate which samples are in the same branch (e.g. with torch.where(S==i), given i is a branch number). However, after passing them through a branch together, you would have to restore their original order in the batch. Maybe you could try something like:</p>
<pre><code class="lang-python">    def forward(self, X: torch.Tensor, S: torch.Tensor) -&gt; torch.Tensor:
        output = torch.zeros(X.shape[0], self.out_features)
        for i in range(self.num_branches):
            s = torch.where(S==i)[0]
            if s.shape[0] &gt; 0:
                output[s] = self.output_fcs[i](X[s])
        return output
</code></pre>
<p>Test it out and see whichever works and suits your use case more : )</p>

flpgrz: <p>Thanks a lot, it looks promising and I’ll test it tomorrow. In the forward you index X using variable i, but I don’t see where i comes from <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

ID56: <p>You are right:</p>
<pre><code class="lang-python">return torch.cat([self.output_fcs[s](X[i, :]).unsqueeze(0) for s in S], axis=0)
</code></pre>
<p>should be</p>
<pre><code class="lang-python">return torch.cat([self.output_fcs[s](X[i, :]).unsqueeze(0) for i, s in enumerate(S)], axis=0)
</code></pre>
<p>I edited into the original post too. Thanks!</p>

flpgrz: <p>I think this solution works - can we imagine also a way without a for loop?</p>

ID56: <aside class="quote no-group" data-username="flpgrz" data-post="5" data-topic="151978">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/flpgrz/48/44557_2.png" class="avatar"> flpgrz:</div>
<blockquote>
<p>can we imagine also a way without a for loop?</p>
</blockquote>
</aside>
<p>Probably. For that we would have to think like the second solution I suggested: compute the solution of each branch sequentially. Actually, we can try to exploit <code>nn.Sequential</code> for this, as sequential is really equivalent to iterating over a list of layers.</p>
<pre data-code-wrap="python"><code class="lang-python">class LinearBranch(nn.Module):
    def __init__(self, in_features: int, out_features: int, branch_id: int):
        super().__init__()
        self.branch_id = branch_id
        self.linear = nn.Linear(in_features, out_features)

    def forward(self, inputs):
        """inputs: Tuple of (X, S, output)."""
        s = torch.where(inputs[1]==self.branch_id)[0]
        if s.shape[0] &gt; 0:
            inputs[2][s] = self.linear(inputs[0][s])
        return inputs

class MyModel(nn.Module):
    "A model with variable number of branches."
    def __init__(self, in_features: int = 16, out_features: int = 1, num_branches: int = 5) -&gt; None:
        super().__init__()
        self.output_fcs = nn.Sequential(*[LinearBranch(in_features, out_features, id) for id in range(num_branches)])
        self.num_branches = num_branches
        self.in_features = in_features
        self.out_features = out_features

    def forward(self, X: torch.Tensor, S: torch.Tensor) -&gt; torch.Tensor:
        output = torch.zeros(X.shape[0], self.out_features)
        _, _, output = self.output_fcs((X, S, output))
        return output
</code></pre>
<p>We can put a wrapper around nn.Linear to handle the checking of the selector. Sequential would then pass X, S and an output tensor through the branches one by one, and each time some samples in the batch would be passed through a branch linear (or none at all), depending on the selector.</p>
<p>Just to check:</p>
<pre data-code-wrap="python"><code class="lang-python">batch_size = 4
num_branches = 5

model = MyModel(16, 1, num_branches)

dummy_input = torch.randn(batch_size, 16)
dummy_selector = torch.randint(0, num_branches, (batch_size,))

print(dummy_selector)
# tensor([2, 1, 2, 4])

output = model(dummy_input, dummy_selector)
print(output, output.shape, output.requires_grad)
# tensor([[ 0.1223],
#         [-0.0652],
#         [ 0.2669],
#         [-0.2918]], grad_fn=&lt;IndexPutBackward0&gt;) torch.Size([4, 1]) True
</code></pre>
<p>Seems like it works, but do let me know how it goes.</p>

ID56: <p>Even if we store all the weight matrices together, we would still have to 1) extract/slice appropriate branch weight, then 2) compute the matrix multiplication of that branch separately – as long as the distribution of the inputs to branches is not uniform (ie for a given batch, a branch can have 0 samples and will be inactive, and two branches a and b can have different number of assigned samples).</p>
<p>It would also not be possible with just a wrapper around linear; to store a 3D weight matrix, we would need to implement a custom layer with custom nn.Parameter() weight and bias, custom reset_parameters() method, etc. which is too much extra work.</p>
<p>And lastly, torch constructs its computation graph dynamically every pass any way, so there is actually no cost of change.</p>
<p>In the special case that the input has the same / uniform number of samples per branch (like say, In a batch of 10 inputs, it is guaranteed that there are 2 samples for each of 5 branches), it may be possible to compute the linear operation very efficiently as a single opearation. Something like this:</p>
<pre><code class="lang-python">import torch

num_branch, samples_per_branch = 5, 2
in_features, out_features = 16, 1

X = torch.randn(num_branch, samples_per_branch, in_features)
# X -&gt; (5, 2, 16)
W = torch.randn(num_branch, in_features, out_features)
# W -&gt; (5, 16, 1)

# loop doing 5 matmuls
out1 = torch.stack([X[i] @ W[i] for i in range(num_branch)])

# single einsum reduction
out2 = torch.einsum("ijk,ikl-&gt;ijl", X, W)

print(out1.shape, out2.shape, torch.allclose(out1, out2))
# torch.Size([5, 2, 1]) torch.Size([5, 2, 1]) True
</code></pre>
<p>But this scenario is very different from the main idea of this question, and thus not very relevant.</p>

pmachart: <p>Thanks again <a class="mention" href="/u/id56">@ID56</a> for the inspiration. It is actually to implement this efficient idea for arbitrary batch compositions using einsum:</p>
<pre><code class="lang-auto">import torch

num_branch = 5
batch_size = 8
in_features, out_features = 16, 1

X = torch.randn(batch_size, in_features)
W = torch.randn(num_branch, in_features, out_features)

# the vector containing the branch indices for each sample
X_branch_idx = np.random.randint(0, high=n_branches, size=batch_size)

# one may stack the W slices corresponding
Wstack = torch.stack([W[i,:,:] for i in X_branch_idx])

# and compute the output for the whole batch using einsum
out = torch.einsum("ij,ijk-&gt;ik", X, Wstack)
</code></pre>
<p>There’s still a for loop to build Wstack but that should be very little overhead.</p>

ID56: <p>Glad I could be useful! Note that for large batches this may grow exponentially in memory. For example, if you have batch size = 512, you will need to allocate memory for a tensor of size (512, 16, 1) during each forward pass. Regardless, it is still a nice solution for when the batch (and model weights) isn’t too large.</p>