PyTorch torch.nn.Module Function

PyTorch torch.nn 参考手册PyTorch torch.nn Reference Manual

torch.nn.ModuleIt is the base class for all neural network modules in PyTorch.

All custom network models should inherit from this class, as it provides parameter management, device migration, model saving, and other functions.

Class Definition

torch.nn.Module

Main Attributes

  • parameters(): Returns all learnable parameters of the model.
  • named_parameters(): Returns an iterator of parameter names and values.
  • children(): Returns the model's submodules.
  • named_children(): Returns an iterator of submodule names and modules.
  • modules(): Returns all modules.
  • state_dict(): Returns a dictionary containing all parameters.

Usage Examples

Example 1: Creating a Custom Module

Inherit nn.Module to create a custom network:

Example

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(SimpleNet, self).__init__()
        # Define network layers
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        # Define forward propagation
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# Instantiate the model
model = SimpleNet(input_dim=784, hidden_dim=256, output_dim=10)

# Test
x = torch.randn(32, 784)
output = model(x)

print(Model structure:)
print(model)
print(nInput shape:, x.shape)
print(Output shape:, output.shape)

Example 2: Managing Parameters

Access and manage model parameters:

Example

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3)
        self.bn1 = nn.BatchNorm2d(16)
        self.fc = nn.Linear(16 * 14 * 14, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = Net()

# Count the number of parameters
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

print(Total parameters:, total_params)
print(Trainable parameters:, trainable_params)

# Access specific parameters
print(nconv1 weight shape:, model.conv1.weight.shape)
print(fc bias shape:, model.fc.bias.shape)

Example 3: Model Saving and Loading

Save and load the model:

Example

import torch
import torch.nn as nn
import tempfile
import os

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 5)

    def forward(self, x):
        return self.fc(x)

model = Net()

# Save the entire model
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
    torch.save(model, f.name)
    path_full = f.name

# Save state_dict (recommended method)
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
    torch.save(model.state_dict(), f.name)
    path_state = f.name

# Load the model
loaded_model = Net()
loaded_model.load_state_dict(torch.load(path_state))
loaded_model.eval()

# Test the loaded model
x = torch.randn(2, 10)
output1 = model(x)
output2 = loaded_model(x)

print(Original output:, output1[0].tolist())
print(Output after loading:, output2[0].tolist())

# Clean up
os.remove(path_full)
os.remove(path_state)

Example 4: Device Migration

Migrate the model between different devices:

Example

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(10, 5)

    def forward(self, x):
        return self.fc(x)

model = Net()

# Check the current device
print(Parameter device:, model.fc.weight.device)

# Move to GPU (if available)
if torch.cuda.is_available():
    model = model.cuda()
    print(After moving to GPU:, model.fc.weight.device)

# Move back to CPU
model = model.cpu()
print(After moving back to CPU:, model.fc.weight.device)

Example 5: Initializing with apply

Use apply to initialize recursively:

Example

import torch
import torch.nn as nn

def init_weights(module):
    if isinstance(module, nn.Linear):
        nn.init.xavier_uniform_(module.weight)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Conv2d):
        nn.init.kaiming_normal_(module.weight)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(10, 20)
        self.fc2 = nn.Linear(20, 5)

    def forward(self, x):
        return self.fc2(self.fc1(x))

model = Net()
model.apply(init_weights)

print(fc1 weight:, model.fc1.weight[0, :3].tolist())
print(fc2 weight:, model.fc2.weight[0, :3].tolist())

Example 6: Complex Network Structure

Build a network with branches:

Example

import torch
import torch.nn as nn

class BranchNet(nn.Module):
    def __init__(self):
        super(BranchNet, self).__init__()
        # Backbone
        self.shared = nn.Linear(10, 20)
        # Branch
        self.branch_a = nn.Linear(20, 5)
        self.branch_b = nn.Linear(20, 3)

    def forward(self, x):
        feat = self.shared(x)
        out_a = self.branch_a(feat)
        out_b = self.branch_b(feat)
        return out_a, out_b

model = BranchNet()
x = torch.randn(4, 10)
out_a, out_b = model(x)

print(Branch A output:, out_a.shape)
print(Branch B output:, out_b.shape)

FAQ

Q1: Why must super().__init__() be called?

It calls the parent class's initialization method, ensuring that the module registers parameters correctly.

<h3 Q2: How to view the model structure?

Directly print(model) or use the summary function from torchvision.

<h3 Q3: What is the difference between parameters() and modules()?

parameters() returns learnable parameters, while modules() returns all modules.

Use Cases

nn.ModuleIt is the foundation for building all custom neural networks. The main application scenarios include:

  • Defining network structure: Any custom neural network
  • Model management: Saving, loading, migration
  • Parameter management: Optimizer updates, gradient computation

Note: All custom Modules must implement the forward method.

PyTorch torch.nn 参考手册PyTorch torch.nn Reference Manual

Other Extensions