Named layers in tensorboard graph of dynamic Module

username: hdkgr
cooked: <p>When I directly define a nn.Module with fixed layers and use <code>torch.utils.tensorboard.SummaryWriter</code>, I get a nice vizualisation of my computational graph like this::</p>
<pre><code class="lang-auto">
class NeuralNetStrategy(nn.Module):
    def __init__(self, input_length):
        super(NeuralNetStrategy, self).__init__()
        self.fc1 = nn.Linear(input_length, 10) 
        self.fc_out = nn.Linear(10, 1)
        
    def forward(self, x):
        x = self.fc1(x).tanh()
        x = self.fc_out(x).relu()
        return x


model = NeuralNetStrategy(1)
with SummaryWriter(logdir) as writer:
    writer.add_graph(some_input_tensor) 
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/a/a1b1c07db894a3beab891b59897e544b5a89b5c7.png" data-download-href="https://discuss.pytorch.org/uploads/default/a1b1c07db894a3beab891b59897e544b5a89b5c7" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/a/a1b1c07db894a3beab891b59897e544b5a89b5c7_2_254x250.png" alt="image" data-base62-sha1="n4pHFZ9ZNvIGPs2xvtar7VprhvF" width="254" height="250" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/a/a1b1c07db894a3beab891b59897e544b5a89b5c7_2_254x250.png, https://discuss.pytorch.org/uploads/default/optimized/2X/a/a1b1c07db894a3beab891b59897e544b5a89b5c7_2_381x375.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/2X/a/a1b1c07db894a3beab891b59897e544b5a89b5c7_2_508x500.png 2x" data-dominant-color="EDE5E6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">612×601 19.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>My goal is to get a similar output graph but for a dynamic number of layers created at runtime, but I’m having trouble to make this work.</p>
<p>I’ve tried using <code>nn.ModuleDict()</code> or building an OrderedDict and passing it to <code>nn.Sequential</code> in order to create named layers, but my tensorboard graph does not display the names. My current class definition and tensorboard output looks something like this:</p>
<pre><code class="lang-auto">class NeuralNetStrategy(nn.Module):
    def __init__(self, input_length, hidden_nodes: Iterable[int], hidden_activations: Iterable):
        super(NeuralNetStrategy, self).__init__()
        
        self.layers = nn.ModuleDict() 
        self.layers['fc_0'] = nn.Linear(input_length, hidden_nodes[0])
        self.layers['activation_0'] = hidden_activations[0]

        for i in range (1, len(hidden_nodes)):
            self.layers['fc_' + str(i)] = nn.Linear(hidden_nodes[i-1], hidden_nodes[i])
            self.layers['activation_' + str(i)] = hidden_activations[i]

        self.layers['fc_out'] = nn.Linear(hidden_nodes[-1], 1)
        self.layers['activation_out'] = nn.ReLU()
        
    def forward(self, x):
        for layer in self.layers.values():
            x = layer(x)
        return x
</code></pre>
<p>When I initialize my model, it looks fine in pytorch (and does perform the desired computations), yet the tensorboard graph does not respect the named layers:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/2X/f/f1d55f780a5c7282e1a8701878a94d26ad5b347f.png" alt="image" data-base62-sha1="yvmdCLzQEFX7mOydfKpaDE9HXn9" width="314" height="171"></p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/2X/3/3d574454f399f9d95fb2fef11fc7c452f53cbe96.png" alt="image" data-base62-sha1="8KE5Q9OdfUmDsqwmV9khp8v80FU" width="306" height="184"></p>
<p>How can I fix this?</p>