Function/Class for Generating Variable-Depth Networks

username: covfefe
cooked: <p>I am working on a project which includes training many neural nets of the same “format” but varying depth. It is thus useful to have a function/class that takes “depth” as an argument–or more generally a variable-length list of layer sizes–and returns a neural net of the requisite depth. For example, in the case of simple feedforward ReLU nets, the following function suffices:</p>
<pre><code class="lang-auto">def make_relu_net(layer_sizes):
    stack = []
    for i, input_size in enumerate(layer_sizes[:-1]):
        output_size = layer_sizes[i+1]
        stack.append(nn.Linear(input_size, output_size))
        if i &lt; len(layer_sizes) - 2:
            stack.append(nn.ReLU())
    return nn.Sequential(*stack)
</code></pre>
<p>However, I am wondering how to extend this to more general neural net layouts that don’t necessarily fit the <code>nn.Sequential</code> format. It’s my understanding that <code>nn.Module</code> should be used for more general network construction. Given the above example of ReLU nets for simplicity, my naive attempt would look like this:</p>
<pre><code class="lang-auto">class FeedForwardNet(nn.Module):
    def __init__(self, layer_sizes):
        super().__init__()
        self.stack = []
        for i, input_size in enumerate(layer_sizes[:-1]):
            output_size = layer_sizes[i+1]
            self.stack.append(
                nn.Linear(input_size, output_size))

    def forward(self, x):
        for layer in self.stack[:-1]:
            x = layer(x).clamp(min=0)
        return self.stack[-1](x)
</code></pre>
<p>There are a few things that seem problematic to me with the above class:</p>
<ol>
<li>
<p>The main issue is that there’s a for-loop in the <code>forward</code> method. In all the pytorch code examples of deep nets I’ve seen online, I don’t think I’ve seen any with for-loops in <code>forward</code>. My assumption is because this would be slow compared to hard-coding the network layout. But my use-case is to allow for variable-depth, so I’m not sure how to get around this. Is it okay to have the for-loop there or is there a better way, e.g., not subclassing <code>nn.Module</code>?</p>
</li>
<li>
<p>Given that all the layers are stored in <code>self.stack</code>, it seems the <code>state_dict</code> is never populated. Is there a way to generate a proper <code>state_dict</code> that still allows for a variable number of layers as required in my use-case?</p>
</li>
<li>
<p>Smaller issue: is the proper thing to include ReLU layers in <code>__init__</code> as attributes of the class (like the <code>Linear</code> layers) or using <code>clamp</code> in <code>forward</code> as in my code? More generally for any non-weight layer: better to compute on-the-fly in <code>forward</code> or make an attribute in <code>__init__</code>?</p>
</li>
</ol>

username: ptrblck
cooked: <ol>
<li>
<p>The for loop is not a problem and you shouldn’t see any Python overhead, as the main workload is in the layer execution, not the for loop evaluation in Python. <a href="https://pytorch.org/docs/stable/nn.html#torch.nn.ModuleList" rel="nofollow noopener"><code>nn.ModuleList</code></a> uses a for loop as the example use case.</p>
</li>
<li>
<p>If you store the modules in a Python list, they won’t be registered properly, and <code>model.parameters()</code> won’t return all parameters. Use <code>nn.ModuleList</code> instead.</p>
</li>
<li>
<p>It depends on your coding style. Stateless modules, e.g. <code>ReLU</code> can be used in their functional API (<code>F.relu</code>) without the necessity to store some parameters manually. If you register these modules in your <code>__init__</code>, you could e.g. swap them with another activation function without manipulating the <code>forward</code> method, but I would still say it’s more a coding style question.</p>
</li>
</ol>