Revision history [back]

To obtain all the layers of a PyTorch model, you can use the children() method of the model. This method returns an iterator over all the modules in the model.

Here's an example:

import torch

# Define your PyTorch model
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
        self.conv2 = torch.nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.fc = torch.nn.Linear(128 * 16 * 16, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = x.view(-1, 128 * 16 * 16)
        x = self.fc(x)
        return x

# Create an instance of the model
model = MyModel()

# Get all the layers using the `children()` method
for layer in model.children():
    print(layer)

This will output:

Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
Linear(in_features=32768, out_features=10, bias=True)