How to create Stacked network?

Tanay_Rastogi: <p>I am bit puzzled on how can I create a model where there is a block is just repeated multiple times and are in parallel to each other? The inputs are different to each block.</p>
<p><strong>Task:</strong><br>
I want to create a classification model that should look like the image. There is a Zone block which is repeated multiple times. All these zone blocks have exact same parameters - they are just copies. But they have different inputs. Mostly, the “y_x” input to each of the block. The output of all of them will go to a softmax layer to give the probability to choose one of those block.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/b/3bd4a96501034f1177d4a2c9dc1d8cedccd4416f.png" data-download-href="https://discuss.pytorch.org/uploads/default/3bd4a96501034f1177d4a2c9dc1d8cedccd4416f" title="DNN"><img src="https://discuss.pytorch.org/uploads/default/original/3X/3/b/3bd4a96501034f1177d4a2c9dc1d8cedccd4416f.png" alt="DNN" data-base62-sha1="8xhNn8AhwmICEivhKFEXr62xSVN" width="339" height="500" data-dominant-color="9DBAC7"><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">DNN</span><span class="informations">561×826 9.98 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><strong>Current Solution:</strong><br>
My solution is to use the ModuleList as shown below.</p>
<pre><code class="lang-auto">class ZoneBlock(nn.Module):
    def __init__(self,
                 input_dim:int,
                 num_neurons:int,
                 num_hlayers:int):
        super().__init__()

        #### MODEL #####
        seq = []

        # Input Layers
        seq += [nn.Linear(input_dim, num_neurons, dtype=torch.float32), nn.ReLU()]
        # Hidden Layers
        for n in range(num_hlayers):
            # Second Hidden Layers
            seq += [nn.Linear(num_neurons, num_neurons, dtype=torch.float32), nn.ReLU()]
        # Output layer.
        seq += [nn.Linear(num_neurons, 1, dtype=torch.float32)]

        #### MODEL ####
        self.block = nn.Sequential(*seq)

    def forward(self, data_input):
        """Apply the model to the input"""
        data = self.block(data_input)
        return data
</code></pre>
<p>The I stack the zoneBlock like this.</p>
<pre><code class="lang-auto">class Stacked(nn.Module):
    def __init__(self, n_zones:int, input_dim:int,
                               num_neurons=16, num_hlayers=1):
        super().__init__()
       
        #### Zone Block #####
        zone = ZoneBlock(input_dim= input_dim,
                         num_neurons=num_neurons,
                         num_hlayers=num_hlayers)      
        
        #### MODEL ####
        self.zoneSet = nn.ModuleList([zone for z in range(n_zones)])
    
    def forward(self, X):
        # Shape of X: [n_ind x n_zone  x n_attrb]

        # List for taking the logits for each zone
        zone_outputs = list()

        # Reshaping the input to [n_zone x n_ind x n_attrb] to loop eaisly
        x_reshaped = torch.reshape(X, (X.shape[1], X.shape[0], X.shape[2]))

        # Loop across all zones
        for idx, dataIn in enumerate(x_reshaped):
            zone_outputs.append(torch.flatten(self.zoneSet[idx](dataIn)))
        
        ## Return
        # Output Shape: [n_ind x n_zone]
        output = torch.stack(zone_outputs).transpose(0, 1)       
        return output

</code></pre>
<p><strong>Issue:</strong><br>
Is there better way to do it? I think because I am using a for loop in the forward, the training is quite slow there.</p>
<p>Thanks!</p>

sofiakim: <p>Consider using nn.ModuleList for efficient stacking of identical blocks. To improve training speed, you can parallelize computation using torch.cat instead of a loop in the forward pass.</p>