How to create an efficient conditional layer in pytorch?

username: bartley
cooked: <p>I have a resnet50 model that outputs a class prediction (1, 2 or 3). Based on the output of the classifier, I want to make another prediction that selects the next layer/model based on the previous model output.</p>
<p>This is what I have so far.</p>
<pre><code class="lang-python">import torch

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.model1 = torch.nn.Linear(1, 1, bias=False)
        torch.nn.init.ones_(self.model1.weight)

        self.model2 = torch.nn.Linear(1, 1, bias=False)
        torch.nn.init.ones_(self.model2.weight)

        self.model3 = torch.nn.Linear(1, 1, bias=False)
        torch.nn.init.ones_(self.model3.weight)

    def forward(self, x):
        
        # Get batch_size
        batch_size = x.size(1)
        output = torch.zeros(batch_size, 1, device=x.device)
        
        # Loop over every value in batch
        for i in range(batch_size):
            value = x[:, i]
            if value == 1:
                output[i] = self.model1(value)
            elif value == 2:
                output[i] = self.model2(value)
            else:
                output[i] = self.model3(value)

        return output
</code></pre>
<pre><code class="lang-auto">model = SimpleModel()

output = model(torch.tensor([[1,2,3]], dtype=torch.float32))
output

tensor([[1.],
        [2.],
        [3.]], grad_fn=&lt;CopySlices&gt;)
</code></pre>
<p>My concern is that I am only computing one forward pass on each iteration of the loop which seems very inefficient. What happens if I increase the batch size to 64? Will the forward pass be computed in parallel?</p>
<p>Any thoughts/ideas would be appreciated.</p>

username: eqy
cooked: <p>It depends on the size of the model, but some simple initial solutions you could try are:</p>
<ol>
<li>simply running all three models for each input and only selecting the chosen input (e.g., via gather <a href="https://pytorch.org/docs/stable/generated/torch.gather.html" class="inline-onebox" rel="noopener nofollow ugc">torch.gather — PyTorch 2.0 documentation</a>)</li>
<li>accumulating input for each model until there is enough to run an entire batch (for whatever batch size you choose)</li>
</ol>
<p>Beyond that I think efficiently data-dependent control flow at a per-input basis is a tricky problem.<br>
There also has been interesting work in leveraging sparse primitives for Mixture-of-Experts training (MoE) for LLM <a href="https://arxiv.org/pdf/2211.15841.pdf" rel="noopener nofollow ugc">https://arxiv.org/pdf/2211.15841.pdf</a>, so I wonder if that could be applicable here.</p>
<p>If you’re okay with e.g., “soft-conditioning,” then this paper might be interesting: <a href="https://arxiv.org/abs/1904.04971" class="inline-onebox" rel="noopener nofollow ugc">[1904.04971] CondConv: Conditionally Parameterized Convolutions for Efficient Inference</a></p>

username: bartley
cooked: <p>Hey <a class="mention" href="/u/eqy">@eqy</a>, thanks for the reply.</p>
<p>I really like the idea of the second solution. Even though it would require multiple iterations over the dataset, I think it would be the fastest solution. I will try this out!</p>