Iterating the Layers of a Neural Network

Occasionally you may want to iterate through the layers of the network to see the outputs of each layer, or to examine its trained weights. The sample code below demonstrates iterating through a network.

Sample Code

import torch import torch.nn as nn network = nn.Sequential( nn.Linear(2, 5), nn.Tanh(), nn.Linear(5,2) ) x = torch.randn(1, 2) # Dictionary to store the intermediate outputs layer_outputs = {} current_input = x for i, layer in enumerate(network): # Pass the data sequentially through each layer current_input = layer(current_input) # Save the output of the current layer layer_outputs[f"layer_{i}_{layer.__class__.__name__}"] = current_input # Print results for layer_name, output in layer_outputs.items(): print(f"{layer_name} output shape: {output.shape}")