Running multiple Modules in parallel

MasterScrat: <p>I am implementing a multi-head network. (This is to implement multi-head DQN, a specific reinforcement learning method, but this doesn’t really matter here.)</p>
<p>My network has the following architecture:</p>
<pre><code>input -&gt; 128x (separate fully connected layers) -&gt; output averaging
</code></pre>
<p>I am using a ModuleList to hold the list of fully connected layers. Here’s how it looks at this point:</p>
<pre><code>class MultiHead(nn.Module):
    def __init__(self, dim_state, dim_action, hidden_size=32, nb_heads=1):
        super(MultiHead, self).__init__()
        
        self.networks = nn.ModuleList()
        for _ in range(nb_heads):
            network = nn.Sequential(
                nn.Linear(dim_state, hidden_size),
                nn.Tanh(),
                nn.Linear(hidden_size, dim_action)
            )
            self.networks.append(network)
        
        self.cuda()
        self.optimizer = optim.Adam(self.parameters())
</code></pre>
<p>Then, when I need to calculate the output, I use a <code>for ... in</code> construct to perform the forward and backward pass through all the layers:</p>
<pre><code>q_values = torch.cat([net(observations) for net in self.networks])

# skipped code which ultimately computes the loss I need

self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
</code></pre>
<p><strong>This works!</strong> But I am wondering if I couldn’t do this more efficiently. By doing a <code>for...in</code>, I am actually going through each separate FC layer one by one, and as a result the training time grows with the number of FC layers.</p>
<p>Can this operation could be done in parallel?</p>
<p>This is similar to this (unanswered) question: <a href="https://discuss.pytorch.org/t/parallel-execution-of-modules-in-nn-modulelist/43940" class="inline-onebox">Parallel execution of modules in nn.ModuleList</a></p>

MasterScrat: <p>I have tried with CUDA streams but I still see a big slowdown when scaling up the number of heads:</p>
<pre><code>    streams = [torch.cuda.Stream() for _ in range(nb_heads)]
    torch.cuda.synchronize()
    
    losses = []
    net_idx = 0
    for net in self.networks:
        with torch.cuda.stream(streams[net_idx]):
            q_values = net(observations)
            q_values = q_values.gather(1, actions.unsqueeze(1)).squeeze(1)

            next_q_values = net(next_observations)
            next_q_values = next_q_values.max(1)[0]

            expected_q_values = rewards + gamma * next_q_values * (1 - terminals)
            losses.append((q_values - expected_q_values.detach()))
        
        net_idx += 1
    
    torch.cuda.synchronize()
    loss = torch.cat(losses).pow(2).mean()
    
    self.optimizer.zero_grad()
    loss.backward()
    self.optimizer.step()</code></pre>

MasterScrat: <p>I have also tried to use a pool of CUDA streams to avoid having to recreate them at each timestep, and I have removed the <code>torch.cuda.synchronize()</code>s which don’t seem to be useful.</p>
<p>But still, the training speed decreases linearly with the number of heads <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=9" title=":frowning:" class="emoji" alt=":frowning:"></p>

MasterScrat: <p><a href="https://stackoverflow.com/a/58389075/318557" rel="nofollow noopener">Another option</a> may be to use (abuse?) a <code>Conv1d </code> layer, using the <code>groups</code> argument to handle the parallel network simultaneously, and using a kernel size of <code>1</code> to get it to act as a a <code>Linear</code> layer.</p>
<p>Would that be equivalent? Should I expect a speedup from this approach?</p>

vitchyr: <p>I posted this <a href="https://discuss.pytorch.org/t/parallel-execution-of-modules-in-nn-modulelist/43940/7">in the other discussion thread</a>, but it seems that the <a href="https://stackoverflow.com/a/58389075/318557" rel="nofollow noopener">other option</a> that <a class="mention" href="/u/masterscrat">@MasterScrat</a> mentioned is in fact faster if and only if you have a GPU.</p>