Splitting a Pytorch Neural Network
ONce you have trained a network, you may wish to split the network. That is, you split the network between two layers in order to create the input model, and the output model. This is often done when creating auto encoders.Splitting a network can be easily accomplished using Python indexing utilities:
subset = network[:2]
This code creates a new model from the first two layers of the model named "network".
Sample Code
import torch
import torch.nn as nn
network = nn.Sequential(
nn.Linear(2, 5),
nn.Tanh(),
nn.Linear(5,2)
)
#splitting the network into two models
# 1. Split into two sub-models
# Model 1 contains layers up to the Tanh activation (indices 0 and 1)
model_part1 = network[:2]
# Model 2 contains the final linear layer (index 2 onwards)
model_part2 = network[2:]
# 2. Test the split models with a dummy input
x = torch.randn(1, 2)
# Pass through Part 1
hidden_features = model_part1(x)
print("Part 1 Output Shape:", hidden_features.shape) # Expected: [1, 5]
# Pass the intermediate result through Part 2
final_output = model_part2(hidden_features)
print("Part 2 Output Shape:", final_output.shape)