Applying independent nn modules in segregated layers (parallelised)

richardbaxter: <p>I recommend PyTorch provide a method to apply multiple independent nn modules in parallel, such as in a segregated layer. At a minimum I recommend they provide a wrapper function for the nn.Conv1d workaround used to parallelise the execution of independent nn.Linear layers. For example;</p>
<pre><code class="lang-auto">segregatedLinear = LinearSegregated(in_features=in_features, out_features=out_features, number_sublayers=linearSublayersNumber)

class LinearSegregated(nn.Module):
	def __init__(self, in_features, out_features, number_sublayers):
		super().__init__()
		self.segregatedLinear = nn.Conv1d(in_channels=in_features*number_sublayers, out_channels=out_features*number_sublayers, kernel_size=1, groups=number_sublayers)
		self.number_sublayers = number_sublayers
		
	def forward(self, x):
		#x.shape = batch_size, number_sublayers, in_features
		x = x.view(x.shape[0], x.shape[1]*x.shape[2], 1)
		x = self.segregatedLinear(x)
		x = x.view(x.shape[0], self.number_sublayers, x.shape[1]//self.number_sublayers)
		#x.shape = batch_size, number_sublayers, out_features
		return x
</code></pre>

ptrblck: <p><code>nn.Linear</code> layers already accept inputs with additional dimensions as <code>[batch_size, *, in_features]</code> and will apply the layer on each sample in these additional dimensions.<br>
I would assume <code>nn.Conv1d</code> would provide the same functionality with its <code>groups</code> argument.<br>
If you have a good abstraction in mind, could you create a feature request on GitHub so that the code owners could discuss it, please?</p>

richardbaxter: <p>I would assume that <code>nn.Linear</code> would apply (multiply) the same parameters for each <code>x</code> subsample (in <code>dim=1</code>; <code>number_sublayers</code>). The example I provided is based on the following workaround;</p>
<aside class="onebox stackexchange" data-onebox-src="https://stackoverflow.com/questions/58374980/run-multiple-models-of-an-ensemble-in-parallel-with-pytorch/58389075#58389075">
  <header class="source">

      <a href="https://stackoverflow.com/questions/58374980/run-multiple-models-of-an-ensemble-in-parallel-with-pytorch/58389075#58389075" target="_blank" rel="noopener nofollow ugc">stackoverflow.com</a>
  </header>

  <article class="onebox-body">
      <a href="https://stackoverflow.com/users/2913106/flawr" target="_blank" rel="noopener nofollow ugc">
    <img alt="flawr" src="https://i.stack.imgur.com/cM7Yq.jpg?s=256&amp;g=1" class="thumbnail onebox-avatar" width="256" height="256">
  </a>

<h4>
  <a href="https://stackoverflow.com/questions/58374980/run-multiple-models-of-an-ensemble-in-parallel-with-pytorch/58389075#58389075" target="_blank" rel="noopener nofollow ugc">Run multiple models of an ensemble in parallel with PyTorch</a>
</h4>

<div class="tags">
  <strong>python, deep-learning, pytorch, ensemble-learning</strong>
</div>

<div class="date">
  
  answered by
  <a href="https://stackoverflow.com/users/2913106/flawr" target="_blank" rel="noopener nofollow ugc">
    flawr
  </a>
  on <a href="https://stackoverflow.com/questions/58374980/run-multiple-models-of-an-ensemble-in-parallel-with-pytorch/58389075#58389075" target="_blank" rel="noopener nofollow ugc">07:09AM - 15 Oct 19 UTC</a>
</div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


ptrblck: <p>Yes as seen here:</p>
<pre><code class="lang-python">lin = nn.Linear(10, 10)
x = torch.randn(2, 10, 10)
ref = lin(x)

outs = []
for idx in range(x.size(1)):
    x_ = x[:, idx]
    out = lin(x_)
    outs.append(out)
    
outs = torch.stack(outs, dim=1)
print((outs - ref).abs().max())
# tensor(0., grad_fn=&lt;MaxBackward1&gt;)
</code></pre>

richardbaxter: <p>Thank you for your support</p>